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

# Charger le modèle BERT-like (CamemBERT pour l'analyse d'intention)
camembert_model_name = "camembert-base"
camembert_tokenizer = AutoTokenizer.from_pretrained(camembert_model_name)
camembert_model = AutoModelForSequenceClassification.from_pretrained(camembert_model_name)

# Charger le modèle génératif GPT-2 francisé
gpt2_model_name = "dbddv01/gpt2-french-small"
gpt2_tokenizer = AutoTokenizer.from_pretrained(gpt2_model_name)
gpt2_model = AutoModelForCausalLM.from_pretrained(gpt2_model_name)

# Dictionnaire d'intentions de base (vous pouvez l'ajuster pour plus de sophistication)
intent_dict = {
    0: "salutation",
    1: "question_faq",
    2: "aide"
}

# Fonction de gestion des intentions et génération des réponses
def chatbot(user_input):
    # 1. Utilisation de CamemBERT pour analyser l'intention
    cam_inputs = camembert_tokenizer(user_input, return_tensors="pt", max_length=128, truncation=True)
    cam_outputs = camembert_model(**cam_inputs)
    
    # Détection de l'intention prédominante
    intent = torch.argmax(cam_outputs.logits, dim=1).item()
    detected_intent = intent_dict.get(intent, "inconnu")

    # 2. Générer la réponse en fonction de l'intention détectée
    if detected_intent == "salutation":
        response = "Bonjour! Comment puis-je vous aider aujourd'hui ?"
    elif detected_intent == "question_faq":
        # On utilise GPT-2 pour générer une réponse à une question
        gpt2_inputs = gpt2_tokenizer(user_input, return_tensors="pt")
        gpt2_outputs = gpt2_model.generate(gpt2_inputs["input_ids"], max_length=100, pad_token_id=gpt2_tokenizer.eos_token_id)
        response = gpt2_tokenizer.decode(gpt2_outputs[0], skip_special_tokens=True)
    elif detected_intent == "aide":
        response = "Je suis là pour vous aider ! Que puis-je faire pour vous ?"
    else:
        response = "Je ne suis pas sûr de comprendre. Pouvez-vous reformuler votre question ?"

    return response

# Créer une interface Gradio pour tester le chatbot
demo = gr.Interface(
    fn=chatbot, 
    inputs="text", 
    outputs="text", 
    title="Chatbot BERT-GPT en Français"
)

# Lancer l'application Gradio
demo.launch()