Update main.py

main.py

@@ -6,15 +6,12 @@ import redis
 from transformers import (
     AutoTokenizer,
     AutoModelForSequenceClassification,
-    TrainingArguments,
 )
 import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader, Dataset
 from torch.optim import AdamW
 from fastapi import FastAPI, HTTPException, Request
-from pydantic import BaseModel
-from typing import List, Dict
 from fastapi.responses import HTMLResponse
 import multiprocessing
 import time
@@ -29,75 +26,21 @@ REDIS_PASSWORD = os.getenv('REDIS_PASSWORD')
 
 app = FastAPI()
 
-# Configuración de idioma
-language_responses = {
-    "es": {
-        0: [
-            "Lo siento, no entiendo.",
-            "No estoy seguro de entender lo que quieres decir.",
-            "¿Podrías reformular tu pregunta?"
-        ],
-        1: [
-            "Hola! ¿Cómo estás?",
-            "¡Hola! ¿Qué tal?",
-            "Buenos días/tardes/noches, ¿cómo te va?"
-        ],
-        2: [
-            "¿Cómo te puedo ayudar?",
-            "¿En qué puedo ayudarte?",
-            "Dime, ¿qué necesitas?"
-        ],
-        # ... más respuestas para otras clases
-    },
-    "en": {
-        0: [
-            "Sorry, I don't understand.",
-            "I'm not sure I understand what you mean.",
-            "Could you rephrase your question?"
-        ],
-        1: [
-            "Hello! How are you?",
-            "Hi! What's up?",
-            "Good morning/afternoon/evening, how are you doing?"
-        ],
-        2: [
-            "How can I help you?",
-            "What can I do for you?",
-            "Tell me, what do you need?"
-        ],
-        # ... más respuestas para otras clases
-    }
-}
-
-default_language = "es"  # Idioma predeterminado
-
-# Servicio de Chatbot 
+default_language = "es"
+
 class ChatbotService:
     def get_response(self, user_id, message, predicted_class, language=default_language):
-        # Accede al diccionario de respuestas según el idioma
-        responses = language_responses.get(language, language_responses["es"])
-
-        # Lógica basada en la clase predicha
-        if predicted_class == 1:
-            # Seleccionar una respuesta de saludo aleatoria
-            return random.choice(responses[1])
-        elif predicted_class == 2:
-            # Seleccionar una respuesta de ayuda aleatoria
-            return random.choice(responses[2])
-        else:
-            # Seleccionar una respuesta de no comprensión aleatoria
-            return random.choice(responses[0]) 
+        return "Respuesta por defecto." 
 
 chatbot_service = ChatbotService()
 
-# Modelo de clasificación de texto
 class UnifiedModel(nn.Module):
     def __init__(self, models):
         super(UnifiedModel, self).__init__()
         self.models = nn.ModuleList(models)
         hidden_size = self.models[0].config.hidden_size
-        self.projection = nn.Linear(len(models) * 3, 768)  # Dimensión de salida corregida
-        self.classifier = nn.Linear(hidden_size, 3) # 3 clases
+        self.projection = nn.Linear(len(models) * 3, 768)
+        self.classifier = nn.Linear(hidden_size, 3)
 
     def forward(self, input_ids, attention_mask):
         hidden_states = []
@@ -106,11 +49,11 @@ class UnifiedModel(nn.Module):
                 input_ids=input_id,
                 attention_mask=attn_mask
             )
-            hidden_states.append(outputs.logits)  # Usar directamente outputs.logits
+            hidden_states.append(outputs.logits)
 
         concatenated_hidden_states = torch.cat(hidden_states, dim=1)
-        projected_features = self.projection(concatenated_hidden_states) # Proyectar a hidden_size
-        logits = self.classifier(projected_features) 
+        projected_features = self.projection(concatenated_hidden_states)
+        logits = self.classifier(projected_features)
         return logits
 
     @staticmethod
@@ -118,14 +61,12 @@ class UnifiedModel(nn.Module):
         model_name = "unified_model"
         model_data_bytes = redis_client.get(f"model:{model_name}")
         if model_data_bytes:
-            model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3) # 3 clases
+            model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
             model.load_state_dict(torch.load(model_data_bytes))
         else:
-            model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3) # 3 clases
-        
-        return UnifiedModel([model, model]) # Asegurar que se usa una lista de modelos, en este caso 2
+            model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
+        return UnifiedModel([model, model])
 
-# Dataset para entrenamiento
 class SyntheticDataset(Dataset):
     def __init__(self, tokenizers, data):
         self.tokenizers = tokenizers
@@ -146,14 +87,13 @@ class SyntheticDataset(Dataset):
         tokenized["labels"] = torch.tensor(label)
         return tokenized
 
-# Manejo de la conversación
 conversation_history = {}
 
 @app.post("/process")
 async def process(request: Request):
     data = await request.json()
     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
-    
+
     tokenizers = {}
     models = {}
 
@@ -164,10 +104,10 @@ async def process(request: Request):
     tokenizer_data_bytes = redis_client.get(f"tokenizer:{tokenizer_name}")
 
     if model_data_bytes:
-        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3) # 3 clases
+        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
         model.load_state_dict(torch.load(model_data_bytes))
     else:
-        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3) # 3 clases
+        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
     models[model_name] = model
 
     if tokenizer_data_bytes:
@@ -179,53 +119,49 @@ async def process(request: Request):
 
     unified_model = UnifiedModel(list(models.values()))
     unified_model.to(torch.device("cpu"))
-    
+
     if data.get("train"):
         user_data = data.get("user_data", [])
         if not user_data:
             user_data = [
-                {"text": "Hola", "label": 1}, 
+                {"text": "Hola", "label": 1},
                 {"text": "Necesito ayuda", "label": 2},
                 {"text": "No entiendo", "label": 0}
-                # ... más ejemplos para otras clases
             ]
-        
+
         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)
 
-        # Memoria de Conversación
         if user_id not in conversation_history:
             conversation_history[user_id] = []
         conversation_history[user_id].append(text)
 
-        # Concatenar el historial al mensaje actual (puedes usar otra técnica)
-        contextualized_text = " ".join(conversation_history[user_id][-3:]) # Usar los últimos 3 mensajes
+        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 = get_chatbot_response(user_id, text, predicted_class, language)
         return {"answer": response}
-    
+
     else:
         raise HTTPException(status_code=400, detail="Request must contain 'train' or 'message'.")
 
 def get_chatbot_response(user_id, question, predicted_class, language):
-    # Almacenar el mensaje en el historial
     if user_id not in conversation_history:
         conversation_history[user_id] = []
     conversation_history[user_id].append(question)
@@ -401,7 +337,7 @@ def push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name):
 
 def continuous_training():
     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
-    
+
     while True:
         try:
             data = redis_client.lpop("training_queue")
@@ -409,12 +345,12 @@ def continuous_training():
                 data = json.loads(data)
                 unified_model = UnifiedModel.load_model_from_redis(redis_client)
                 unified_model.train()
-                
+
                 train_dataset = SyntheticDataset(data["tokenizers"], data["data"])
                 train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
-                
+
                 optimizer = AdamW(unified_model.parameters(), lr=5e-5)
-                
+
                 for epoch in range(10):
                     for batch in train_loader:
                         input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in data["tokenizers"].keys()]
@@ -425,9 +361,9 @@ def continuous_training():
                         loss.backward()
                         optimizer.step()
                         optimizer.zero_grad()
-                        
+
                         print(f"Epoch {epoch}, Loss {loss.item()}")
-                
+
                 push_to_redis(unified_model.models, data["tokenizers"], redis_client, "unified_model", "unified_tokenizer")
             time.sleep(10)
         except Exception as e: