Update agent.py

agent.py

@@ -1,17 +1,13 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer, pipeline
-import torch
+from transformers import pipeline
 
-class HuggingFaceAgent:
-    def __init__(self, model_id, access_token):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id, token=access_token)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            token=access_token,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto"
-        )
-        self.pipeline = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer)
+# Usaremos un modelo ligero para mejor compatibilidad
+modelo = "google/gemma-2b-it"
 
-    def responder(self, prompt):
-        respuesta = self.pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]["generated_text"]
-        return respuesta[len(prompt):].strip()
+generador = pipeline("text-generation", model=modelo)
+
+def ejecutar_agente(texto):
+    """
+    Procesa una pregunta usando el modelo LLM.
+    """
+    resultado = generador(texto, max_new_tokens=300, do_sample=True, temperature=0.7)
+    return resultado[0]["generated_text"]