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

# Charger le modèle LLaMA 3.2 et le tokenizer
model_name = "meta-llama/LLaMA-3.2"  # Remplacez par le bon nom de modèle si nécessaire
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Définir une fonction pour générer une réponse
def generate_response(input_text):
    # Encoder le texte d'entrée
    inputs = tokenizer(input_text, return_tensors="pt")

    # Générer la réponse
    with torch.no_grad():
        outputs = model.generate(**inputs, max_length=150)  # Ajustez max_length selon vos besoins

    # Décoder et retourner la réponse
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Créer l'interface Gradio
interface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Entrez votre texte ici..."),
    outputs="text",
    title="Générateur de texte avec LLaMA 3.2",
    description="Entrez un texte et obtenez une réponse générée par le modèle LLaMA 3.2."
)

# Lancer l'application
if __name__ == "__main__":
    interface.launch()