Orbii-KM / app.py
Pgohari's picture
Update app.py
ed9832f verified
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()