Update main.py

main.py

@@ -6,6 +6,7 @@ import redis
 from transformers import (
     AutoTokenizer,
     AutoModelForSequenceClassification,
+    AutoModelForCausalLM,
 )
 import torch
 import torch.nn as nn
@@ -29,8 +30,46 @@ app = FastAPI()
 default_language = "es"
 
 class ChatbotService:
-    def get_response(self, user_id, message, predicted_class, language=default_language):
-        return "Respuesta por defecto." 
+    def __init__(self):
+        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
+        self.model_name = "response_model"
+        self.tokenizer_name = "response_tokenizer"
+
+    def get_response(self, user_id, message, language=default_language):
+        model = self.load_model_from_redis()
+        tokenizer = self.load_tokenizer_from_redis()
+
+        if model is None or tokenizer is None:
+            return "El modelo aún no está listo. Por favor, inténtelo de nuevo más tarde."
+
+        input_text = f"Usuario: {message} Asistente:"
+        input_ids = tokenizer.encode(input_text, return_tensors="pt").to("cpu")
+
+        with torch.no_grad():
+            output = model.generate(input_ids=input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True)
+
+        response = tokenizer.decode(output[0], skip_special_tokens=True)
+        response = response.replace(input_text, "").strip()
+
+        return response
+
+    def load_model_from_redis(self):
+        model_data_bytes = self.redis_client.get(f"model:{self.model_name}")
+        if model_data_bytes:
+            model = AutoModelForCausalLM.from_pretrained("gpt2")
+            model.load_state_dict(torch.load(model_data_bytes))
+            return model
+        else:
+            return None
+
+    def load_tokenizer_from_redis(self):
+        tokenizer_data_bytes = self.redis_client.get(f"tokenizer:{self.tokenizer_name}")
+        if tokenizer_data_bytes:
+            tokenizer = AutoTokenizer.from_pretrained("gpt2")
+            tokenizer.add_tokens(json.loads(tokenizer_data_bytes))
+            return tokenizer
+        else:
+            return None
 
 chatbot_service = ChatbotService()
 
@@ -155,7 +194,13 @@ async def process(request: Request):
             logits = unified_model(input_ids=input_ids, attention_mask=attention_mask)
             predicted_class = torch.argmax(logits, dim=-1).item()
 
-        response = get_chatbot_response(user_id, text, predicted_class, language)
+        response = chatbot_service.get_response(user_id, contextualized_text, language)
+
+        redis_client.rpush("training_queue", json.dumps({
+            "tokenizers": {tokenizer_name: tokenizer.get_vocab()},
+            "data": [{"text": contextualized_text, "label": predicted_class}]
+        }))
+
         return {"answer": response}
 
     else:
@@ -166,7 +211,7 @@ def get_chatbot_response(user_id, question, predicted_class, language):
         conversation_history[user_id] = []
     conversation_history[user_id].append(question)
 
-    return chatbot_service.get_response(user_id, question, predicted_class, language)
+    return chatbot_service.get_response(user_id, question, language)
 
 @app.get("/")
 async def get_home():
@@ -364,7 +409,13 @@ def continuous_training():
 
                         print(f"Epoch {epoch}, Loss {loss.item()}")
 
-                push_to_redis(unified_model.models, data["tokenizers"], redis_client, "unified_model", "unified_tokenizer")
+                push_to_redis(
+                    {"response_model": unified_model},
+                    {"response_tokenizer": tokenizer},
+                    redis_client,
+                    "response_model",
+                    "response_tokenizer",
+                )
             time.sleep(10)
         except Exception as e:
             print(f"Error in continuous training: {e}")