medical-reasoning-gpt-oss-20b

#pip install torch --index-url https://download.pytorch.org/whl/cu128
#pip install "trl>=0.20.0" "peft>=0.17.0" "transformers>=4.55.0"

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

base_model_name = "openai/gpt-oss-20b"
adapter_name = "dousery/medical-reasoning-gpt-oss-20b"

tokenizer = AutoTokenizer.from_pretrained(base_model_name)

base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

model = PeftModel.from_pretrained(base_model, adapter_name)
model = model.merge_and_unload()

messages = [
    {"role": "system", "content": "You are a medical reasoning assistant."},
    {"role": "user", "content": (
        """A 55-year-old man has chest pain and elevated troponin I without ST elevation.
         What is the diagnosis and what additional test would you order next?"""
    )}
]

prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=2048,
    temperature=0.2,
    do_sample=False
)

raw_output = tokenizer.decode(outputs[0], skip_special_tokens=False)

#  PARSING THE OUTPUT
thinking_pattern = r"<\|end\|><\|start\|>assistant<\|channel\|>analysis<\|message\|>(.*?)<\|end\|>"
final_pattern = r"<\|start\|>assistant<\|channel\|>final<\|message\|>(.*?)<\|return\|>"

thinking_match = re.search(thinking_pattern, raw_output, re.DOTALL)
final_match = re.search(final_pattern, raw_output, re.DOTALL)

thinking_text = thinking_match.group(1).strip() if thinking_match else "N/A"
final_text = final_match.group(1).strip() if final_match else "N/A"

print("Thinking:", thinking_text)
print("\nFinal:", final_text)