Update agent.py

agent.py

@@ -1,63 +1,17 @@
-import agent_tools
-from smolagents import CodeAgent, AzureOpenAIServerModel
-import my_tokens
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer, pipeline
+import torch
 
-# Instancia del modelo
-model = AzureOpenAIServerModel(
-    model_id=my_tokens.AZURE_OPENAI_MODEL,
-    azure_endpoint=my_tokens.AZURE_OPENAI_ENDPOINT,
-    api_key=my_tokens.AZURE_OPENAI_API_KEY,
-    api_version=my_tokens.OPENAI_API_VERSION    
-)
-
-class CustomAgent:
-    def __init__(self):
-        # Definir agentes de trabajo con sus herramientas
-        self.web_agent = CodeAgent(
-            model=model,
-            tools=[agent_tools.google_search, agent_tools.wiki_search, agent_tools.visit_page, agent_tools.final_answer],
-            max_steps=8,
-            name="web_agent",
-            description="Este agente maneja búsquedas web."
-        )
-
-        self.audio_agent = CodeAgent(
-            model=model,
-            tools=[agent_tools.speech_to_text_tool, agent_tools.final_answer],
-            max_steps=4,
-            name="audio_agent",
-            description="Este agente convierte audio a texto."
-        )
-
-        self.py_agent = CodeAgent(
-            model=model,
-            tools=[agent_tools.do_python, agent_tools.final_answer],
-            additional_authorized_imports=["json", "pandas", "numpy", "regex"],
-            max_steps=8,
-            name="python_code_agent",
-            description="Este agente ejecuta y valida código Python."
-        )
-
-        self.visual_agent = CodeAgent(
-            model=model,
-            tools=[agent_tools.visual_qa_tool, agent_tools.final_answer],
-            max_steps=4,
-            name="visual_qa_agent",
-            description="Este agente responde preguntas sobre imágenes."
-        )
-
-        self.manager_agent = CodeAgent(
-            model=model,
-            tools=[],
-            managed_agents=[self.web_agent, self.audio_agent, self.py_agent, self.visual_agent],
-            planning_interval=8,
-            verbosity_level=2,
-            max_steps=12,
+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)
 
-    def forward(self, question: str, attachment: str = None) -> str:
-        if attachment:
-            result = self.manager_agent.run(question, additional_args={"attachment": attachment})
-        else:
-            result = self.manager_agent.run(question)
-        return result
+    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()