import gradio as gr
import torch
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer

# --- Configuration des Modèles ---
model_id = "NaA-IA/melta-english"
translation_id = "Helsinki-NLP/opus-mt-en-fr"

print("Chargement du modèle de texte... (cela peut prendre quelques minutes)")
# Chargement du tokenizer et du modèle de langue
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

print("Chargement du traducteur...")
# Pipeline de traduction
translator = pipeline("translation_en_to_fr", model=translation_id)

def generate_response(message, history):
    """
    Fonction de génération : prend le message, génère en anglais, traduit en français.
    """
    # 1. Préparation de l'entrée (en anglais)
    inputs = tokenizer(message, return_tensors="pt").to(model.device)
    
    # 2. Génération de la réponse brute
    with torch.no_grad():
        output_tokens = model.generate(
            **inputs, 
            max_new_tokens=150,
            do_sample=True,
            temperature=0.7,
            top_p=0.9
        )
    
    # 3. Décodage de la réponse (on retire le prompt d'origine)
    full_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
    response_en = full_text[len(message):].strip()
    
    # 4. Traduction vers le français
    if response_en:
        translation = translator(response_en)
        response_fr = translation[0]['translation_text']
    else:
        response_fr = "Désolé, je n'ai pas pu générer de réponse."

    return response_fr

# --- Interface Gradio ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"# 🤖 Chat Melta-English (Traduit en FR)\nModèle : `{model_id}`")
    
    chatbot = gr.ChatInterface(
        fn=generate_response,
        examples=["Tell me a story about a dragon", "What is the capital of France?"],
        title="Assistant de Traduction Automatique"
    )

if __name__ == "__main__":
    demo.launch()