Update app.py

app.py

@@ -8,6 +8,10 @@ import csv
 import json
 import os
 import torch
+from langchain.document_loaders import JSONLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
 
 huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
 
@@ -15,7 +19,7 @@ huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
 if huggingface_token:
     login(token=huggingface_token)
 
-# Configuración del modelo
+# Configuración del modelo de generación de texto
 @st.cache_resource
 def load_llm():
     llm = HuggingFaceEndpoint(
@@ -39,6 +43,32 @@ classification_model, classification_tokenizer = load_classification_model()
 
 id2label = {0: "multas", 1: "politicas_de_privacidad", 2: "contratos", 3: "denuncias", 4: "otros"}
 
+# Cargar documentos JSON para cada categoría
+@st.cache_resource
+def load_json_documents():
+    documents = {}
+    categories = ["multas", "politicas_de_privacidad", "contratos", "denuncias", "otros"]
+    for category in categories:
+        with open(f"./{category}.json", "r", encoding="utf-8") as f:
+            data = json.load(f)["questions_and_answers"]
+            documents[category] = [entry["question"] + " " + entry["answer"] for entry in data]
+    return documents
+
+json_documents = load_json_documents()
+
+# Configuración de Embeddings y Vector Stores
+@st.cache_resource
+def create_vector_store():
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-l6-v2", model_kwargs={"device": "cpu"})
+    vector_stores = {}
+    for category, docs in json_documents.items():
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
+        split_docs = [doc for doc in text_splitter.split_text(docs)]
+        vector_stores[category] = FAISS.from_texts(split_docs, embeddings)
+    return vector_stores
+
+vector_stores = create_vector_store()
+
 def classify_text(text):
     inputs = classification_tokenizer(text, return_tensors="pt", max_length=4096, truncation=True, padding="max_length")
     classification_model.eval()
@@ -47,7 +77,7 @@ def classify_text(text):
     logits = outputs.logits
     predicted_class_id = logits.argmax(dim=-1).item()
     predicted_label = id2label[predicted_class_id]
-    return f"Clasificación: {predicted_label}\n\nDocumento:\n{text}"
+    return predicted_label
 
 def translate(text, target_language):
     template = '''
@@ -105,61 +135,62 @@ def handle_uploaded_file(uploaded_file):
         return str(e)
 
 def main():
-    st.title("LexAIcon: Traduce, resume y explica textos legales")
-  
-    st.header("Puedes conversar con este chatbot basado en Mistral7B-Instruct y subir archivos para que ser traducidos resumidos o explicados.")
-
     st.image("./icon.jpg", width=100)
+    st.title("LexAIcon")
+    st.write("Puedes conversar con este chatbot basado en Mistral7B-Instruct y subir archivos para que el chatbot los procese.")
 
+    if "messages" not in st.session_state:
+        st.session_state["messages"] = [{"role": "assistant", "content": "¿Cómo puedo ayudarte?"}]
 
-    if "generated" not in st.session_state:
-        st.session_state["generated"] = []
+    with st.sidebar:
+        st.text_input("HuggingFace Token", value=huggingface_token, type="password", key="huggingface_token")
+        st.caption("[Consigue un HuggingFace Token](https://huggingface.co/settings/tokens)")
 
-    # Entrada del usuario
-    user_input = st.text_input("Tú: ", "")
+    for msg in st.session_state.messages:
+        st.chat_message(msg["role"]).write(msg["content"])
 
-    # Botones de Resumir, Traducir y Explicar
-    operation = st.radio("Selecciona una operación", ["Resumir", "Traducir", "Explicar"])
+    if prompt := st.chat_input():
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        st.chat_message("user").write(prompt)
 
-    target_language = None
-    summary_length = None
+        operation = st.radio("Selecciona una operación", ["Resumir", "Traducir", "Explicar"])
+        target_language = None
+        summary_length = None
 
-    if operation == "Traducir":
-        target_language = st.selectbox("Selecciona el idioma de traducción", ["español", "inglés", "francés", "alemán"])
+        if operation == "Traducir":
+            target_language = st.selectbox("Selecciona el idioma de traducción", ["español", "inglés", "francés", "alemán"])
 
-    if operation == "Resumir":
-        summary_length = st.selectbox("Selecciona la longitud del resumen", ["corto", "medio", "largo"])
+        if operation == "Resumir":
+            summary_length = st.selectbox("Selecciona la longitud del resumen", ["corto", "medio", "largo"])
 
-    # Manejo de archivos subidos
-    uploaded_files = st.file_uploader("Sube un archivo", type=["txt", "pdf", "docx", "csv", "json"], accept_multiple_files=True)
-
-    if st.button("Enviar"):
-        if user_input:
-            response = llm_engine_hf.invoke(user_input)
-            st.session_state.generated.append({"user": user_input, "bot": response['generated_text']})
-
-    if st.button("Ejecutar"):
-        if uploaded_files:
+        if uploaded_files := st.file_uploader("Sube un archivo", type=["txt", "pdf", "docx", "csv", "json"], accept_multiple_files=True):
             for uploaded_file in uploaded_files:
                 file_content = handle_uploaded_file(uploaded_file)
-                if operation == "Resumir":
-                    if summary_length == "corto":
-                        length = "de aproximadamente 50 palabras"
-                    elif summary_length == "medio":
-                        length = "de aproximadamente 100 palabras"
-                    elif summary_length == "largo":
-                        length = "de aproximadamente 500 palabras"
-                    result = summarize(file_content, length)
-                elif operation == "Traducir":
-                    result = translate(file_content, target_language)
-                elif operation == "Explicar":
-                    result = classify_text(file_content)
-                st.write(result)
-
-    if st.session_state.get("generated"):
-        for chat in st.session_state["generated"]:
-            st.write(f"Tú: {chat['user']}")
-            st.write(f"Chatbot: {chat['bot']}")
+                classification = classify_text(file_content)
+                vector_store = vector_stores[classification]
+                search_docs = vector_store.similarity_search(prompt)
+                context = " ".join([doc.page_content for doc in search_docs])
+                prompt_with_context = f"Contexto: {context}\n\nPregunta: {prompt}"
+                response = llm_engine_hf.invoke(prompt_with_context)
+                msg = response.content
+
+        elif operation == "Resumir":
+            if summary_length == "corto":
+                length = "de aproximadamente 50 palabras"
+            elif summary_length == "medio":
+                length = "de aproximadamente 100 palabras"
+            elif summary_length == "largo":
+                length = "de aproximadamente 500 palabras"
+            msg = summarize(prompt, length)
+
+        elif operation == "Traducir":
+            msg = translate(prompt, target_language)
+
+        else:
+            msg = llm_engine_hf.invoke(prompt).content
+
+        st.session_state.messages.append({"role": "assistant", "content": msg})
+        st.chat_message("assistant").write(msg)
 
 if __name__ == "__main__":
     main()