Update main.py

main.py

@@ -18,6 +18,8 @@ from typing import List, Dict
 from fastapi.responses import HTMLResponse
 import multiprocessing
 import time
+import uuid
+import random
 
 load_dotenv()
 
@@ -27,12 +29,74 @@ 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 
+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]) 
+
+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.classifier = nn.Linear(len(models) * hidden_size, 2)
+        self.classifier = nn.Linear(len(models) * hidden_size, 3) # 3 clases
 
     def forward(self, input_ids, attention_mask):
         hidden_states = []
@@ -51,12 +115,13 @@ 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=2)
+            model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3) # 3 clases
             model.load_state_dict(torch.load(model_data_bytes))
         else:
-            model = AutoModelForSequenceClassification.from_pretrained("gpt2")
+            model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3) # 3 clases
         return UnifiedModel([model])
 
+# Dataset para entrenamiento
 class SyntheticDataset(Dataset):
     def __init__(self, tokenizers, data):
         self.tokenizers = tokenizers
@@ -77,22 +142,11 @@ 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):
-    """
-    Processes requests for training and prediction.
-
-    Args:
-        request (Request): The incoming request object.
-
-    Returns:
-        dict: A dictionary containing either a message indicating successful
-              training data submission or the model's prediction.
-
-    Raises:
-        HTTPException: If the request does not contain 'train' or 'predict'
-                       keys.
-    """
     data = await request.json()
     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
     
@@ -106,10 +160,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=2)
+        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3) # 3 clases
         model.load_state_dict(torch.load(model_data_bytes))
     else:
-        model = AutoModelForSequenceClassification.from_pretrained("gpt2")
+        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3) # 3 clases
     models[model_name] = model
 
     if tokenizer_data_bytes:
@@ -125,9 +179,13 @@ async def process(request: Request):
     if data.get("train"):
         user_data = data.get("user_data", [])
         if not user_data:
-            user_data = [{"text": "Sample text for automatic training.", "label": 0}]
+            user_data = [
+                {"text": "Hola", "label": 1}, 
+                {"text": "Necesito ayuda", "label": 2},
+                {"text": "No entiendo", "label": 0}
+                # ... más ejemplos para otras clases
+            ]
         
-        # Add user data to Redis queue for asynchronous training
         redis_client.rpush("training_queue", json.dumps({
             "tokenizers": {tokenizer_name: tokenizer.get_vocab()},
             "data": user_data
@@ -135,9 +193,20 @@ async def process(request: Request):
 
         return {"message": "Training data received. Model will be updated asynchronously."}
     
-    elif data.get("predict"):
-        text = data['text']
-        tokenized_inputs = [tokenizers[name](text, return_tensors="pt") for name in tokenizers.keys()]
+    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
+
+        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]
         
@@ -145,114 +214,178 @@ async def process(request: Request):
             logits = unified_model(input_ids=input_ids, attention_mask=attention_mask)
             predicted_class = torch.argmax(logits, dim=-1).item()
         
-        return {"prediction": predicted_class}
+        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 'predict'.")
-
-@app.post("/external_answer")
-async def external_answer(request: Request):
-    """
-    Provides an answer to a question using the unified model and triggers
-    asynchronous training with the new question-answer pair.
+        raise HTTPException(status_code=400, detail="Request must contain 'train' or 'message'.")
 
-    Args:
-        request (Request): The incoming request object containing the question.
+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)
 
-    Returns:
-        dict: A dictionary containing the answer to the question.
-
-    Raises:
-        HTTPException: If the request does not contain a 'question' key.
-    """
-    data = await request.json()
-    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
-    
-    question = data.get('question')
-    if not question:
-        raise HTTPException(status_code=400, detail="Question is required.")
-
-    # Load the model and tokenizer from Redis
-    unified_model = UnifiedModel.load_model_from_redis(redis_client)
-    unified_model.to(torch.device("cpu"))
-
-    tokenizer_data_bytes = redis_client.get(f"tokenizer:unified_tokenizer")
-    if tokenizer_data_bytes:
-        tokenizer = AutoTokenizer.from_pretrained("gpt2")
-        tokenizer.add_tokens(json.loads(tokenizer_data_bytes))
-    else:
-        tokenizer = AutoTokenizer.from_pretrained("gpt2")
-    
-    tokenized_input = tokenizer(question, return_tensors="pt")
-    input_ids = tokenized_input['input_ids']
-    attention_mask = tokenized_input['attention_mask']
-    
-    with torch.no_grad():
-        logits = unified_model(input_ids=[input_ids], attention_mask=[attention_mask])
-        predicted_class = torch.argmax(logits, dim=-1).item()
-    response = {"answer": f"Response to '{question}' is class {predicted_class}"}
-    
-    # Asynchronously train on the new data point
-    redis_client.rpush("training_queue", json.dumps({
-        "tokenizers": {"unified_tokenizer": tokenizer.get_vocab()},
-        "data": [{"text": question, "label": predicted_class}]
-    }))
-    
-    return response
+    return chatbot_service.get_response(user_id, question, predicted_class, language)
 
 @app.get("/")
 async def get_home():
-    """
-    Serves a basic HTML page as the home route.
-
-    Returns:
-        HTMLResponse: The HTML content of the home page.
-    """
-    html_code = """
+    user_id = str(uuid.uuid4())
+    html_code = f"""
     <!DOCTYPE html>
     <html>
     <head>
         <meta charset="UTF-8">
         <title>Chatbot</title>
         <style>
-            body {
-                font-family: Arial, sans-serif;
+            body {{
+                font-family: 'Arial', sans-serif;
                 background-color: #f4f4f9;
                 margin: 0;
                 padding: 0;
-            }
-            .container {
-                max-width: 1200px;
-                margin: 0 auto;
-                padding: 20px;
-            }
-            h1 {
+                display: flex;
+                align-items: center;
+                justify-content: center;
+                min-height: 100vh;
+            }}
+
+            .container {{
+                background-color: #fff;
+                border-radius: 10px;
+                box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);
+                overflow: hidden;
+                width: 400px;
+                max-width: 90%;
+            }}
+
+            h1 {{
                 color: #333;
-            }
+                text-align: center;
+                padding: 20px;
+                margin: 0;
+                background-color: #f8f9fa;
+                border-bottom: 1px solid #eee;
+            }}
+
+            #chatbox {{
+                height: 400px;
+                padding: 20px;
+                overflow-y: auto;
+            }}
+
+            .message {{
+                margin-bottom: 15px;
+                padding: 10px;
+                border-radius: 5px;
+                max-width: 70%;
+                animation: slide-in 0.3s ease-out;
+            }}
+
+            .user-message {{
+                text-align: right;
+                background-color: #eee;
+                margin-left: 30%;
+            }}
+
+            .bot-message {{
+                text-align: left;
+                background-color: #ccf5ff;
+                margin-right: 30%;
+            }}
+
+            #input-area {{
+                display: flex;
+                padding: 10px;
+                background-color: #f8f9fa;
+                border-top: 1px solid #eee;
+            }}
+
+            #message-input {{
+                flex: 1;
+                padding: 10px;
+                border: 1px solid #ccc;
+                border-radius: 5px;
+                margin-right: 10px;
+            }}
+
+            #send-button {{
+                padding: 10px 15px;
+                background-color: #28a745;
+                color: white;
+                border: none;
+                cursor: pointer;
+                border-radius: 5px;
+                transition: background-color 0.3s ease;
+            }}
+
+            #send-button:hover {{
+                background-color: #218838;
+            }}
+
+            @keyframes slide-in {{
+                from {{
+                    transform: translateX(-100%);
+                    opacity: 0;
+                }}
+                to {{
+                    transform: translateX(0);
+                    opacity: 1;
+                }}
+            }}
         </style>
     </head>
     <body>
         <div class="container">
-            <h1>Chatbot Interface</h1>
+            <h1>Chatbot</h1>
+            <div id="chatbox"></div>
+            <div id="input-area">
+                <input type="hidden" id="user-id" value="{user_id}">
+                <input type="text" id="message-input" placeholder="Escribe tu mensaje...">
+                <button id="send-button">Enviar</button>
+            </div>
         </div>
+        <script>
+            const chatbox = document.getElementById('chatbox');
+            const messageInput = document.getElementById('message-input');
+            const sendButton = document.getElementById('send-button');
+            const userId = document.getElementById('user-id').value;
+
+            sendButton.addEventListener('click', sendMessage);
+
+            function sendMessage() {{
+                const message = messageInput.value;
+                if (message.trim() === '') return;
+
+                appendMessage('user', message);
+                messageInput.value = '';
+
+                fetch('/process', {{
+                    method: 'POST',
+                    headers: {{
+                        'Content-Type': 'application/json'
+                    }},
+                    body: JSON.stringify({{ message: message, user_id: userId, language: 'es' }})
+                }})
+                .then(response => response.json())
+                .then(data => {{
+                    appendMessage('bot', data.answer);
+                }});
+            }}
+
+            function appendMessage(sender, message) {{
+                const messageElement = document.createElement('div');
+                messageElement.classList.add('message', `${{sender}}-message`);
+                messageElement.textContent = message;
+                chatbox.appendChild(messageElement);
+                chatbox.scrollTop = chatbox.scrollHeight;
+            }}
+        </script>
     </body>
     </html>
     """
     return HTMLResponse(content=html_code)
 
 def push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name):
-    """
-    Saves the given models and tokenizers to Redis.
-
-    Args:
-        models (dict): A dictionary of model names and their corresponding
-                       PyTorch models.
-        tokenizers (dict): A dictionary of tokenizer names and their
-                           corresponding tokenizers.
-        redis_client: The Redis client instance.
-        model_name (str): The base name to use for saving the models.
-        tokenizer_name (str): The base name to use for saving the tokenizers.
-    """
     for model_name, model in models.items():
         torch.save(model.state_dict(), model_name)
         redis_client.set(f"model:{model_name}", open(model_name, "rb").read())
@@ -262,9 +395,6 @@ def push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name):
         redis_client.set(f"tokenizer:{tokenizer_name}", json.dumps(tokens))
 
 def continuous_training():
-    """
-    Continuously checks for new training data in Redis and updates the model.
-    """
     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
     
     while True:
@@ -300,10 +430,8 @@ def continuous_training():
             time.sleep(5)
 
 if __name__ == "__main__":
-    # Start the continuous training process in a separate process
     training_process = multiprocessing.Process(target=continuous_training)
     training_process.start()
 
-    # Run the FastAPI app
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)