#V01

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Liste des modèles disponibles
model_list = [
    "fbaldassarri/tiiuae_Falcon3-1B-Instruct-autogptq-int8-gs128-asym",
    "MisterAI/jpacifico_Chocolatine-3B-Instruct-DPO-v1.2",
    # Ajoutez d'autres modèles ici
]

def load_model(model_name):
    """Charge le modèle et le tokenizer"""
    if model_name is not None:
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
        return model, tokenizer
    else:
        return None, None

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)

def main(model_name, input_text, max_length, temperature):
    """Fonction principale pour générer le texte"""
    if model_name is not None:
        model, tokenizer = load_model(model_name)
        generated_text = generate_text(model, tokenizer, input_text, max_length, temperature)
        return generated_text
    else:
        return "Veuillez sélectionner un modèle"

demo = gr.Blocks()

with demo:
    gr.Markdown("# Modèle de Langage")
    
    with gr.Row():
        model_select = gr.Dropdown(model_list, label="Sélectionner un modèle")
    with gr.Row():
        load_button = gr.Button("Charger le modèle")
        
    with gr.Row():
        input_text = gr.Textbox(label="Texte d'entrée")
    with gr.Row():
        max_length_slider = gr.Slider(50, 500, label="Longueur maximale", value=200)
        temperature_slider = gr.Slider(0.1, 1.0, label="Température", value=0.7)
    with gr.Row():
        submit_button = gr.Button("Soumettre")
        
    output_text = gr.Textbox(label="Texte généré")
    history = gr.JSON(label="Historique")
    
    load_button.click(
        load_model,
        inputs=model_select,
        outputs=None,
        queue=False
    )
    
    submit_button.click(
        main,
        inputs=[model_select, input_text, max_length_slider, temperature_slider],
        outputs=output_text,
        queue=False
    )

if __name__ == "__main__":
    demo.launch()