#https://huggingface.co/spaces/MisterAI/Try_Small_Models02
#TSMS_app.py_V09
#Add SystemInfo to complete requirements.txt
#Change Interface
##Necessaire de Restart HFSpace tous les 2 ou 3 modèles pour décharger la RAM -- A Voir 
#def unload_model(model):
#    """Décharge le modèle pour libérer la mémoire"""
#    del model  # Supprime la référence au modèle
#    gc.collect()  # Force la collecte des objets non référencés


import subprocess
import sys, platform
from importlib import metadata as md

#SystemInfo
print("Python:", platform.python_version(), sys.implementation.name)
print("OS:", platform.uname())
print("\n".join(sorted(f"{d.metadata['Name']}=={d.version}" for d in md.distributions())))


# 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 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, history):
    """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)    
    # Mettre à jour l'historique
    updated_history = append_history(history, f"Entrée: {input_text}\nSortie: {generated_text}")   
    # Renvoie le texte généré et l'historique mis à jour
    return generated_text, updated_history  

# 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, history
#    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}"

# 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, history  # Ajoutez history ici pour l'utiliser
    try:
        model, tokenizer = load_model(model_name)
        history = append_history(history, f"Modèle {model_name} chargé avec succès")  # Mettre à jour l'historique
        return history  # Renvoie l'historique mis à jour
    except Exception as e:
        history = append_history(history, f"Erreur lors du chargement du modèle {model_name}: {e}")  # Mettre à jour l'historique en cas d'erreur
        return history  # Renvoie l'historique mis à jour
        
        
        
#Fonction Pour Conserver un Historique
def append_history(history, message):
    return f"{history}\n\n***\n{message}" if history else message
        
# Création de l'interface Gradio
demo = gr.Blocks()

with demo:
    with gr.Row():
        gr.Markdown("# TrySmallModels Quickly!\n"
                    "NB : Be Patient - Models Are Running On HF Free Space - Work Better with Models Under 3B Parameters")

    with gr.Row():
        gr.Markdown("## Sélectionner Un Modèle")
        
    with gr.Row():
        with gr.Column():
            # Textbox pour entrer le nom du modèle
            model_name_textbox = gr.Textbox(label="Nom du modèle", placeholder="Depot/NomModele Exemple kurakurai/Luth-LFM2-1.2B", interactive=True)
        with gr.Column():
            # Bouton pour charger le modèle sélectionné
            LoadModel_button = gr.Button("Charger Le Modèle")            
        with gr.Column():
            # Textbox pour Logs et Historique
            history = gr.Textbox(label="Historique", lines=5)
            
    with gr.Row():
        gr.Markdown("## Entrer Un Prompt")
        
    with gr.Row():
        # Textbox pour entrer le texte d'entrée
        input_text = gr.Textbox(label="Texte d'entrée", lines=5)
        
    with gr.Row():
        # Sliders pour ajuster la longueur maximale et la température
        max_length_slider = gr.Slider(50, 4500, label="Longueur maximale", value=200)
        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():
        gr.Markdown("## Sortie LLM")
    
    with gr.Row():
        # Textbox pour afficher le texte généré
        output_text = gr.Textbox(label="Texte généré", lines=5)

    # 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=[history]
    )
    
    # Lors du clic sur "Soumettre", appeler la fonction main
    submit_button.click(
        main,
        inputs=[input_text, max_length_slider, temperature_slider, model_name_textbox, history],
        outputs=[output_text, history],
        queue=False
    )

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