from unsloth import FastLanguageModel
from transformers import TextStreamer
import torch

# Load the finetuned model
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "papasega/gpt-oss-20b-HF4-Multilingual-Thinking", # Replace with your model name if different
    max_seq_length = 128, # Set to the max_seq_length you want 
    dtype = None, # Use None for auto detection
    load_in_4bit = True, # Set to True if you saved in 4bit
)

# Prepare the input message
messages = [
    {"role": "system", "content": "reasoning language: French\n\nYou are a helpful assistant that can solve mathematical problems."},
    {"role": "user", "content": "Résout cette equation pour un élève en classe de seconde : x^4 + 2 = 0."},
]

inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt = True,
    return_tensors = "pt",
    return_dict = True,
    reasoning_effort = "low", # Choose "low", "medium", or "high"
).to(model.device)

# Generate the response
_ = model.generate(**inputs, max_new_tokens = 128, streamer = TextStreamer(tokenizer))