from transformers import pipeline
import gradio as gr
import os
import spaces

# Charger le modèle GPT de Hugging Face
model_id = "deepseek-ai/deepseek-llm-7b-chat"
pipe = pipeline("text-generation", model=model_id)

# Consigne que le modèle suivra pour chaque chat
system_message = """Classify the text into neutral, negative or positive.
Text: This movie is definitely one of my favorite movies of its kind. The interaction between respectable and morally strong characters is an ode to chivalry and the honor code amongst thieves and policemen.
Sentiment:
"""

@spaces.GPU
# Fonction pour générer une réponse à partir du message de l'utilisateur
def generate_response(user_message, history):
    try:
        # Initialiser history s'il est None
        if history is None:
            history = []

        # Ajouter le message système au début de l'historique des messages
        messages = [{"role": "system", "content": system_message}] + [{"role": "user", "content": user_message}]
        
        # Créer une chaîne de caractères représentant la conversation
        conversation_text = system_message + "\n"
        for past_user_input, past_bot_response in history:
            conversation_text += f"User: {past_user_input}\nAssistant: {past_bot_response}\n"
        
        # Ajouter le nouveau message de l'utilisateur
        conversation_text += f"User: {user_message}\n"
        
        # Générer une réponse
        result = pipe(conversation_text, max_new_tokens=150)
        response = result[0]['generated_text'].split("User: ")[-1].strip()  # Extraire la réponse générée
        
        # Mettre à jour l'historique
        history.append((user_message, response))
        return history, response
    except Exception as e:
        # En cas d'erreur, retourner l'historique inchangé et un message d'erreur
        return history, f"Error: {str(e)}"

# Configurer et lancer l'interface de chat avec Gradio
iface = gr.Interface(
    fn=generate_response,
    inputs=[gr.Textbox(lines=2, placeholder="Enter your message here..."), gr.State()],
    outputs=[gr.State(), gr.Textbox()],
    live=True,
)

iface.launch()