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

titulo = """# 🙋🏻‍♂️Bienvenido al Demo de 🌟Yi-Coder-9B-Chat!"""
descripcion = """Yi-Coder-9B-Chat es un modelo de 9B parámetros ajustado para tareas de codificación. Esta demo muestra su capacidad para generar código basado en tus prompts. Yi-Coder es una serie de modelos de lenguaje de código abierto que ofrece un rendimiento de codificación de vanguardia con menos de 10 mil millones de parámetros. Sobresale en la comprensión de contextos largos con una longitud máxima de contexto de 128K tokens. - Soporta 52 lenguajes de programación principales."""

# Definir el dispositivo y la ruta del modelo
dispositivo = "cuda" if torch.cuda.is_available() else "cpu"
ruta_modelo = "01-ai/Yi-Coder-9B-Chat"

# Cargar el tokenizador y el modelo
tokenizador = AutoTokenizer.from_pretrained(ruta_modelo)
modelo = AutoModelForCausalLM.from_pretrained(ruta_modelo, device_map="auto").eval()

@spaces.GPU(duration=130)
def generar_codigo(prompt_sistema, prompt_usuario, max_longitud):
    mensajes = [
        {"role": "system", "content": prompt_sistema},
        {"role": "user", "content": prompt_usuario}
    ]
    texto = tokenizador.apply_chat_template(
        mensajes,
        tokenize=False,
        add_generation_prompt=True
    )
    entradas_modelo = tokenizador([texto], return_tensors="pt").to(dispositivo)
    ids_generados = modelo.generate(
        entradas_modelo.input_ids,
        max_new_tokens=max_longitud,
        eos_token_id=tokenizador.eos_token_id
    )
    ids_generados = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(entradas_modelo.input_ids, ids_generados)
    ]
    respuesta = tokenizador.batch_decode(ids_generados, skip_special_tokens=True)[0]
    return respuesta

def interfaz_gradio():
    with gr.Blocks() as interfaz:
        gr.Markdown(titulo)
        gr.Markdown(descripcion)
        
        prompt_sistema = gr.Textbox(
            label="☯️Instrucción Yi:",
            value="Eres un asistente de codificación útil. Proporciona ejemplos de código claros y concisos.",
            lines=2
        )
        prompt_usuario = gr.Code(
            label="🤔Pregunta de Codificación",
            value="Escribe un algoritmo de quicksort en Python.",
            language="python",
            lines=15
        )
        codigo_salida = gr.Code(label="☯️Yi-Coder-9B", language='python', lines=20, interactive=True)
        max_longitud_slider = gr.Slider(minimum=1, maximum=1800, value=650, label="Longitud Máxima de Tokens")
        
        boton_generar = gr.Button("Generar Código")
        boton_generar.click(
            generar_codigo,
            inputs=[prompt_sistema, prompt_usuario, max_longitud_slider],
            outputs=codigo_salida
        )
    
    return interfaz

if __name__ == "__main__":
    interfaz = interfaz_gradio()
    interfaz.queue()
    interfaz.launch()