Update app.py

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+from transformers import InferenceClient, AutoTokenizer, AutoModelForSequenceClassification
 from PyPDF2 import PdfReader
 from docx import Document
 import csv
@@ -8,8 +9,7 @@ import torch
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from huggingface_hub import login, InferenceClient
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from huggingface_hub import login
 
 huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
 
@@ -36,32 +36,6 @@ 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 = 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()
@@ -72,6 +46,19 @@ def classify_text(text):
     predicted_label = id2label[predicted_class_id]
     return predicted_label
 
+def load_json_documents(category):
+    with open(f"./{category}.json", "r", encoding="utf-8") as f:
+        data = json.load(f)["questions_and_answers"]
+        documents = [entry["question"] + " " + entry["answer"] for entry in data]
+    return documents
+
+def create_vector_store(docs):
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-l6-v2", model_kwargs={"device": "cpu"})
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
+    split_docs = text_splitter.split_text(docs)
+    vector_store = FAISS.from_texts(split_docs, embeddings)
+    return vector_store
+
 def translate(text, target_language):
     template = f'''
     Por favor, traduzca el siguiente documento al {target_language}:
@@ -157,7 +144,8 @@ def main():
             for uploaded_file in uploaded_files:
                 file_content = handle_uploaded_file(uploaded_file)
                 classification = classify_text(file_content)
-                vector_store = vector_stores[classification]
+                docs = load_json_documents(classification)
+                vector_store = create_vector_store(docs)
                 search_docs = vector_store.similarity_search(user_input)
                 context = " ".join([doc.page_content for doc in search_docs])
                 prompt_with_context = f"Contexto: {context}\n\nPregunta: {user_input}"