import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-20b",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, "wjbmattingly/gpt-oss-20b-coref-resolution-lora")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")

system_prompt = """You are an expert at coreference resolution. Given a text containing pronouns and other referring expressions, rewrite the text replacing all pronouns with the full name of the entity they refer to.

Keep the text otherwise identical - only replace pronouns with the names they refer to."""

messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": "John went to the store. He bought milk."},
]

prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=512, do_sample=False)
response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)

Merging the Adapter (Optional)

If you want to merge the LoRA weights into the base model for faster inference:

python

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-20b",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# Load and merge LoRA adapter
model = PeftModel.from_pretrained(base_model, "wjbmattingly/gpt-oss-20b-coref-resolution-lora")
model = model.merge_and_unload()

# Now model is a regular transformers model with merged weights

With MLX on Mac

For MLX, you'll need to first merge the adapter, then convert:

python

# Step 1: Merge and save (run this once)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-20b",
    torch_dtype=torch.float16,
    device_map="cpu",  # Use CPU to avoid GPU memory issues
)
model = PeftModel.from_pretrained(base_model, "wjbmattingly/gpt-oss-20b-coref-resolution-lora")
model = model.merge_and_unload()

# Save merged model
model.save_pretrained("./gpt-oss-20b-coref-resolution-lora-merged", safe_serialization=True)
tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")
tokenizer.save_pretrained("./gpt-oss-20b-coref-resolution-lora-merged")

Then convert to MLX format:

bash

pip install mlx-lm

# Convert to MLX (with quantization for smaller size)
mlx_lm.convert --hf-path ./gpt-oss-20b-coref-resolution-lora-merged --mlx-path ./gpt-oss-20b-coref-resolution-lora-mlx -q

# Generate text
mlx_lm.generate --model ./gpt-oss-20b-coref-resolution-lora-mlx --prompt "Your text here"

Or use directly in Python with MLX:

python

from mlx_lm import load, generate

model, tokenizer = load("./gpt-oss-20b-coref-resolution-lora-mlx")

system_prompt = """You are an expert at coreference resolution. Given a text containing pronouns, rewrite the text replacing all pronouns with the full name of the entity they refer to."""

# Format prompt according to the model's chat template
prompt = f"<|system|>\n{system_prompt}<|end|>\n<|user|>\nJohn went to the store. He bought milk.<|end|>\n<|assistant|>\n"

response = generate(model, tokenizer, prompt=prompt, max_tokens=512)
print(response)

Base Model

Model: openai/gpt-oss-20b

Training Dataset

Dataset: wjbmattingly/synthetic-coref

License

This adapter inherits the license from the base model.