Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 from huggingface_hub import login
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import gradio as gr
 import torch
 
@@ -9,25 +9,24 @@ hf_token = os.getenv("HF_API_TOKEN")
 login(hf_token)
 
 # Cargar el modelo y el tokenizador
-model_name = "bertin-project/bertin-roberta-base-spanish"
+model_name = "mrm8488/t5-base-finetuned-spanish"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-def classify_text(input_text):
+def generate_text(input_text):
     inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
     with torch.no_grad():
-        outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class_id = torch.argmax(logits, dim=-1).item()
-    return predicted_class_id
+        outputs = model.generate(**inputs, max_length=200, num_beams=4, early_stopping=True)
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
 
 # Crear la interfaz con Gradio
 iface = gr.Interface(
-    fn=classify_text,
+    fn=generate_text,
     inputs="text",
     outputs="text",
-    title="Clasificador en Español con BERTin",
-    description="Interfaz para clasificar texto en español usando el modelo BERTin RoBERTa base."
+    title="Generador de Texto en Español",
+    description="Genera texto en español utilizando un modelo de lenguaje preentrenado."
 )
 
-iface.launch()
+iface.launch()