#TSMS_app.py_V07
#A REPARER - CHARGE PLUS CORRECTEMENT LE MODELE

# Import des bibliothèques nécessaires
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import requests
import torch

# Fonction pour charger le modèle et le tokenizer
def load_model(model_name):
    """Charge le modèle et le tokenizer"""
    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    return model, tokenizer

# Fonction pour mettre à jour le label de la Textbox
def update_output_text_label(model_name):
    """Met à jour le label de la Textbox avec le nom du modèle"""
    return f"Réponse de {model_name}"

# Fonction pour générer du texte en utilisant le modèle
def generate_text(model, tokenizer, input_text, max_length, temperature):
    """Génère du texte en utilisant le modèle"""
    inputs = tokenizer(input_text, return_tensors="pt")
    output = model.generate(**inputs, max_length=max_length, temperature=temperature)
    return tokenizer.decode(output[0], skip_special_tokens=True)

# Fonction principale pour générer le texte
def main(input_text, max_length, temperature, model_name):
    """Fonction principale pour générer le texte"""
    global model, tokenizer
    if model is None or tokenizer is None:
        raise Exception("Modèle non chargé. Veuillez charger un modèle avant de générer du texte.")
    generated_text = generate_text(model, tokenizer, input_text, max_length, temperature)
    return generated_text

# Variables globales pour stocker le modèle et le tokenizer
model = None
tokenizer = None

# Fonction pour charger le modèle lors du clic sur le bouton "Charger Le Modèle"
def load_model_button_clicked(model_name):
    """Charge le modèle et le tokenizer lors du clic sur le bouton"""
    global model, tokenizer
    try:
        model, tokenizer = load_model(model_name)
        return f"Modèle {model_name} chargé avec succès"
    except Exception as e:
        return f"Erreur lors du chargement du modèle {model_name}: {e}"

# Création de l'interface Gradio
demo = gr.Blocks()

with demo:
    gr.Markdown("""
    # Choisir un LLM

    **Renseigner un NomRepo/NomModele ou charger le modèle par défaut**

    **Attention : Inférence locale : Choisir un petit modèle : 1.5B maximum**

    **NB : ce HFSpace a pour but de tester rapidement des petits LLM : pas de format GGUF (et autres formats incompatibles listés ci-dessous)**

    **Formats Incompatibles :**
    - GGUF
    - Quantized Models
    - ONNX
    - Diffusers
    - Text-to-Speech (TTS)
    - Text-to-Image (T2I)
    - Stable Diffusion
    - Custom Code
    """)

    with gr.Row():
        # Textbox pour entrer le nom du modèle
        model_name_textbox = gr.Textbox(label="Renseigner Le Nom Du Modèle", value="prithivMLmods/Triangulum-1B", interactive=True)

    with gr.Row():
        # Bouton pour charger le modèle sélectionné
        LoadModel_button = gr.Button("Charger Le Modèle")
    
    with gr.Row():
        # Textbox pour entrer le texte d'entrée
        input_text = gr.Textbox(label="ChatBox")

    with gr.Row():
        # Sliders pour ajuster la longueur maximale et la température
        max_length_slider = gr.Slider(25, 1048, label="Nombre de Token", value=250)
        temperature_slider = gr.Slider(0.1, 1.0, label="Température", value=0.7)

    with gr.Row():
        # Bouton pour soumettre le texte d'entrée
        submit_button = gr.Button("Soumettre")

    with gr.Row():        
        # Textbox pour afficher le texte généré
        output_text = gr.Textbox(label="Réponse", interactive=False)

    # Ajout des interactions pour les boutons
    # Lors du clic sur "Charger Le Modèle", appeler la fonction load_model_button_clicked
#    LoadModel_button.click(
#        load_model_button_clicked,
#        inputs=model_name_textbox,
#        outputs=gr.Textbox(label="Message de chargement")
#    )

    LoadModel_button.click(
        load_model_button_clicked,
        inputs=model_name_textbox,
        outputs=[gr.Textbox(label="Message de chargement"), model_name_textbox]
    ).then(
        update_output_text_label,
        inputs=model_name_textbox,
        outputs=output_text
    )
    
    # Lors du clic sur "Soumettre", appeler la fonction main
    submit_button.click(
        main,
        inputs=[input_text, max_length_slider, temperature_slider, model_name_textbox],
        outputs=output_text,
        queue=False
    )

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