from transformers import pipeline
import gradio as gr

# Cargar modelo ligero (ej: tinyllama, phi-2, o bertin)
model_name = "bertin-project/bertin-gpt-j-6B"  # Modelo en español
tag_generator = pipeline("text-generation", model=model_name)

def generate_tags(title):
    prompt = f"""
    Genera 3 etiquetas técnicas en español para este título de TFM.
    Formato: tag1, tag2, tag3
    Ejemplo: "machine_learning, análisis_de_datos, optimización"

    Título: {title}
    Etiquetas:"""
    
    result = tag_generator(prompt, max_length=100, do_sample=True, temperature=0.3)
    tags = result[0]['generated_text'].split("Etiquetas:")[-1].strip()
    return [tag.strip() for tag in tags.split(",")][:3]

# Interfaz Gradio
demo = gr.Interface(
    fn=generate_tags,
    inputs=gr.Textbox(label="Título del TFM"),
    outputs=gr.JSON(label="Etiquetas generadas"),
    title="Etiquetador automático de TFMs",
    examples=[["Predicción de demanda energética con redes neuronales"]]
)

demo.launch(server_name="0.0.0.0", server_port=7860)