Update main.py

main.py

@@ -39,16 +39,12 @@ class ChatbotService:
     def get_response(self, user_id, message, language=default_language):
         if self.model is None or self.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 = self.tokenizer.encode(input_text, return_tensors="pt").to("cpu")
-
         with torch.no_grad():
             output = self.model.generate(input_ids=input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True)
-
         response = self.tokenizer.decode(output[0], skip_special_tokens=True)
         response = response.replace(input_text, "").strip()
-
         return response
 
     def load_model_from_redis(self):
@@ -64,6 +60,7 @@ class ChatbotService:
         if tokenizer_data_bytes:
             tokenizer = AutoTokenizer.from_pretrained("gpt2")
             tokenizer.add_tokens(json.loads(tokenizer_data_bytes))
+            tokenizer.pad_token = tokenizer.eos_token
             return tokenizer
         return None
 
@@ -82,7 +79,6 @@ class UnifiedModel(nn.Module):
         for model, input_id, attn_mask in zip(self.models, input_ids, attention_mask):
             outputs = model(input_ids=input_id, attention_mask=attn_mask)
             hidden_states.append(outputs.logits)
-
         concatenated_hidden_states = torch.cat(hidden_states, dim=1)
         projected_features = self.projection(concatenated_hidden_states)
         logits = self.classifier(projected_features)
@@ -145,8 +141,10 @@ async def process(request: Request):
     if tokenizer_data_bytes:
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
         tokenizer.add_tokens(json.loads(tokenizer_data_bytes))
+        tokenizer.pad_token = tokenizer.eos_token
     else:
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        tokenizer.pad_token = tokenizer.eos_token
     tokenizers[tokenizer_name] = tokenizer
 
     unified_model = UnifiedModel(list(models.values()))
@@ -160,42 +158,31 @@ async def process(request: Request):
                 {"text": "Necesito ayuda", "label": 2},
                 {"text": "No entiendo", "label": 0}
             ]
-
         redis_client.rpush("training_queue", json.dumps({
             "tokenizers": {tokenizer_name: tokenizer.get_vocab()},
             "data": user_data
         }))
-
         return {"message": "Training data received. Model will be updated asynchronously."}
-
     elif data.get("message"):
         user_id = data.get("user_id")
         text = data['message']
         language = data.get("language", default_language)
-
         if user_id not in conversation_history:
             conversation_history[user_id] = []
         conversation_history[user_id].append(text)
-
         contextualized_text = " ".join(conversation_history[user_id][-3:])
-
         tokenized_inputs = [tokenizers[name](contextualized_text, return_tensors="pt") for name in tokenizers.keys()]
         input_ids = [tokens['input_ids'] for tokens in tokenized_inputs]
         attention_mask = [tokens['attention_mask'] for tokens in tokenized_inputs]
-
         with torch.no_grad():
             logits = unified_model(input_ids=input_ids, attention_mask=attention_mask)
             predicted_class = torch.argmax(logits, dim=-1).item()
-
         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:
         raise HTTPException(status_code=400, detail="Request must contain 'train' or 'message'.")
 
@@ -249,12 +236,15 @@ async def get_home():
             }}
             .message {{
                 margin-bottom: 10px;
+                padding: 10px;
+                border-radius: 5px;
             }}
-            .user {{
-                color: #007bff;
+            .message.user {{
+                background-color: #e1f5fe;
+                text-align: right;
             }}
-            .bot {{
-                color: #28a745;
+            .message.bot {{
+                background-color: #f1f1f1;
             }}
             #input {{
                 display: flex;
@@ -337,13 +327,15 @@ async def get_home():
     return HTMLResponse(content=html_code)
 
 def train_unified_model():
+    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
     while True:
-        redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
         training_queue = redis_client.lrange("training_queue", 0, -1)
         if training_queue:
             for item in training_queue:
                 item_data = json.loads(item)
                 tokenizers = {name: AutoTokenizer.from_pretrained("gpt2") for name in item_data["tokenizers"]}
+                for tokenizer in tokenizers.values():
+                    tokenizer.pad_token = tokenizer.eos_token
                 data = item_data["data"]
                 dataset = SyntheticDataset(tokenizers, data)
                 dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
@@ -375,4 +367,4 @@ if __name__ == "__main__":
     training_process = multiprocessing.Process(target=train_unified_model)
     training_process.start()
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)