Update modules/chatbot.py

modules/chatbot.py

@@ -1,25 +1,59 @@
-import requests
-import os
-from dotenv import load_dotenv
-
-# Cargar variables de entorno
-load_dotenv()
-
-class Llama2Chatbot:
-    def __init__(self):
-        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):
-        payload = {"inputs": prompt}
-        response = requests.post(self.API_URL, headers=self.headers, json=payload)
-        return response.json()[0]['generated_text']
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+import torch
 
 def initialize_chatbot():
-    return Llama2Chatbot()
+    model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+    tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+    return model, tokenizer
 
-def get_chatbot_response(chatbot, prompt):
-    return chatbot.generate_response(prompt)
+def get_chatbot_response(model, tokenizer, prompt, src_lang):
+    tokenizer.src_lang = src_lang
+    encoded_input = tokenizer(prompt, return_tensors="pt")
+    generated_tokens = model.generate(**encoded_input, max_length=100)
+    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+
+def display_chatbot_interface(lang_code):
+    translations = {
+        'es': {
+            'title': "AIdeaText - Chatbot Multilingüe",
+            'input_placeholder': "Escribe tu mensaje aquí...",
+            'send_button': "Enviar",
+        },
+        'en': {
+            'title': "AIdeaText - Multilingual Chatbot",
+            'input_placeholder': "Type your message here...",
+            'send_button': "Send",
+        },
+        'fr': {
+            'title': "AIdeaText - Chatbot Multilingue",
+            'input_placeholder': "Écrivez votre message ici...",
+            'send_button': "Envoyer",
+        }
+    }
+
+    t = translations[lang_code]
+
+    st.header(t['title'])
+
+    if 'chatbot' not in st.session_state:
+        st.session_state.chatbot, st.session_state.tokenizer = initialize_chatbot()
+
+    if 'messages' not in st.session_state:
+        st.session_state.messages = []
+
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    if prompt := st.chat_input(t['input_placeholder']):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        with st.chat_message("assistant"):
+            response = get_chatbot_response(st.session_state.chatbot, st.session_state.tokenizer, prompt, lang_code)
+            st.markdown(response)
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
+        # Guardar la conversación en la base de datos
+        store_chat_history(st.session_state.username, st.session_state.messages)