import gradio as gr
import requests
import os

API_URL = "https://api-inference.huggingface.co/models/gpt2"
api_token = os.environ.get("TOKEN")
headers = {"Authorization": f"Bearer {api_token}"}

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

def generate_yes_no_response(question):
    # Créez un prompt qui force le modèle à choisir entre Oui et Non
    prompt = f"Question : {question}\nRéponse (Oui ou Non) :"
    
    # Configurez les paramètres pour forcer une réponse courte
    payload = {
        "inputs": prompt,
        "parameters": {
            "max_new_tokens": 1,
            "return_full_text": False,
            "do_sample": True,
            "top_k": 2,
            "temperature": 0.7
        }
    }
    
    response = query(payload)
    
    # Assurez-vous que la réponse est soit "Oui" soit "Non"
    if isinstance(response, list) and len(response) > 0:
        answer = response[0].get('generated_text', '').strip().lower()
        if answer == "oui":
            return "Oui"
        elif answer == "non":
            return "Non"
    
    # Si la réponse n'est pas "Oui" ou "Non", retournez une réponse par défaut
    return "Réponse invalide"

def chatbot(message, history):
    response = generate_yes_no_response(message)
    return response

iface = gr.ChatInterface(
    fn=chatbot,
    title="Chatbot Oui/Non",
    description="Posez une question, et je répondrai par Oui ou Non."
)

iface.launch()