Update modules/chatbot.py

modules/chatbot.py

@@ -1,18 +1,22 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+import requests
+import os
+from dotenv import load_dotenv
+
+# Cargar variables de entorno
+load_dotenv()
 
 class Llama2Chatbot:
     def __init__(self):
-        self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
-        self.model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model.to(self.device)
+        self.API_URL = "https://api-inference.huggingface.co/models/meta-llama/Llama-2-7b-hf"
+        api_key = os.getenv("HF_API_KEY")
+        if not api_key:
+            raise ValueError("No se encontró la clave de API de Hugging Face. Asegúrate de configurar la variable de entorno HF_API_KEY.")
+        self.headers = {"Authorization": f"Bearer {api_key}"}
 
-    def generate_response(self, prompt, max_length=100):
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
-        outputs = self.model.generate(**inputs, max_length=max_length)
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response
+    def generate_response(self, prompt):
+        payload = {"inputs": prompt}
+        response = requests.post(self.API_URL, headers=self.headers, json=payload)
+        return response.json()[0]['generated_text']
 
 def initialize_chatbot():
     return Llama2Chatbot()