from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
import torch

MODEL_NAME = "papasega/gpt-oss-20b-mxfp4-HF4-Multilingual-Thinking"

print("🔄 Chargement du modèle (cela peut prendre quelques petites minutes)...\n")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    dtype="auto",
    device_map="cuda",
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

print("✅ Modèle chargé avec succès !")
if torch.cuda.is_available():
    print(f"📊 Mémoire GPU utilisée : {torch.cuda.memory_allocated() / 1e9:.2f} Go")

###----*---#### Génération d'une réponse pour la résolution de l'equation x^4 + 2 = 0.

def generate_response(messages, reasoning_effort="low", max_tokens=512, verbose=True):
    """
    Fonction helper pour générer une réponse

    Args:
        messages (list): Liste de dictionnaires {role, content}
        reasoning_effort (str): "low", "medium", ou "high"
        max_tokens (int): Nombre max de tokens à générer
        verbose (bool): Afficher les détails
    """
    if verbose:
        print(f"🧠 Niveau de raisonnement: {reasoning_effort.upper()}")
        print(f"📝 Génération de {max_tokens} tokens maximum\n")
        print("-" * 70)

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
        reasoning_effort=reasoning_effort,
    ).to(model.device)

    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    with torch.inference_mode():
        _ = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            streamer=streamer,
            temperature=0.7, ###----*---#### Agis dans la créativité du modèle
            top_p=0.9,
            do_sample=True,

        )

    print("\n" + "-" * 70)

messages_exemple1 = [
    {"role": "system", "content": "reasoning language: French\n\nTu es un assistant pédagogique."},
    {"role": "user", "content": "Résout cette equation pour un élève en classe de seconde qui ne connait pas les complexes et élève en classe de Terminale : x^4 + 2 = 0."}
]

generate_response(messages_exemple1, reasoning_effort="low", max_tokens=512)