import gradio as gr
from transformers import pipeline

# Cargar el modelo Latxa
def load_model():
    generator = pipeline("text-generation", model="HiTZ/latxa-7b-v1.1")  # Forzar uso de CPU
    return generator

generator = load_model()

# Función para generar texto
def generate_text(input_text):
    result = generator(input_text, max_new_tokens=50, num_beams=5)
    return result[0]['generated_text']

# Crear la interfaz con Gradio
iface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    title="Latxa",
    description="Idatzi testu bat euskaraz eta Latxa sortuko du jarraipena",
)

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