Update app.py

app.py

@@ -1,83 +1,60 @@
 import streamlit as st
-from huggingface_hub import snapshot_download
-from pathlib import Path
-from mistral_inference.model import Transformer
-from mistral_inference.generate import generate
-from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
-from mistral_common.protocol.instruct.messages import UserMessage
-from mistral_common.protocol.instruct.request import ChatCompletionRequest
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, AutoConfig, AutoModelForSequenceClassification
+from langchain.llms import HuggingFacePipeline
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
+from langchain.embeddings.huggingface import HuggingFaceEmbeddings
 from PyPDF2 import PdfReader
 from docx import Document
 import csv
 import json
-import os
 import torch
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 
-# Descargar y configurar el modelo
-mistral_models_path = Path.home().joinpath('mistral_models', '7B-Instruct-v0.3')
-mistral_models_path.mkdir(parents=True, exist_ok=True)
-snapshot_download(repo_id="mistralai/Mistral-7B-Instruct-v0.3", allow_patterns=["params.json", "consolidated.safetensors", "tokenizer.model.v3"], local_dir=mistral_models_path)
+# Configurar modelo y tokenizador
+model_name = 'mistralai/Mistral-7B-Instruct-v0.3'
+model_config = AutoConfig.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
 
-# Configurar el modelo y el tokenizador
-tokenizer = MistralTokenizer.from_file(f"{mistral_models_path}/tokenizer.model.v3")
-model = Transformer.from_folder(mistral_models_path)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Configuración del modelo de clasificación
-@st.cache_resource
-def load_classification_model():
-    tokenizer_cls = AutoTokenizer.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
-    model_cls = AutoModelForSequenceClassification.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
-    return model_cls, tokenizer_cls
+text_generation_pipeline = pipeline(
+    model=model,
+    tokenizer=tokenizer,
+    task="text-generation",
+    temperature=0.2,
+    repetition_penalty=1.1,
+    return_full_text=True,
+    max_new_tokens=300,
+)
 
-classification_model, classification_tokenizer = load_classification_model()
+prompt_template = """
+### [INST] 
+Instruction: Answer the question based on your knowledge. Here is context to help:
 
-id2label = {0: "multas", 1: "politicas_de_privacidad", 2: "contratos", 3: "denuncias", 4: "otros"}
+{context}
 
-def classify_text(text):
-    inputs = classification_tokenizer(text, return_tensors="pt", max_length=4096, truncation=True, padding="max_length")
-    classification_model.eval()
-    with torch.no_grad():
-        outputs = classification_model(**inputs)
-    logits = outputs.logits
-    predicted_class_id = logits.argmax(dim=-1).item()
-    predicted_label = id2label[predicted_class_id]
-    return predicted_label
+### QUESTION:
+{question} 
 
-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
+[/INST]
+"""
 
-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
+mistral_llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
 
-def translate(text, target_language):
-    completion_request = ChatCompletionRequest(
-        messages=[UserMessage(content=f"Por favor, traduzca el siguiente documento al {target_language}:\n{text}\nAsegúrese de que la traducción sea precisa y conserve el significado original del documento.")]
-    )
-    tokens = tokenizer.encode_chat_completion(completion_request).tokens
-    out_tokens, _ = generate([tokens], model, max_tokens=512, temperature=0.0, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
-    translated_text = tokenizer.instruct_tokenizer.tokenizer.decode(out_tokens[0])
-    return translated_text
+# Crear el prompt desde la plantilla de prompt
+prompt = PromptTemplate(
+    input_variables=["context", "question"],
+    template=prompt_template,
+)
 
-def summarize(text, length):
-    completion_request = ChatCompletionRequest(
-        messages=[UserMessage(content=f"Por favor, haga un resumen {length} del siguiente documento:\n{text}\nAsegúrese de que el resumen sea conciso y conserve el significado original del documento.")]
-    )
-    tokens = tokenizer.encode_chat_completion(completion_request).tokens
-    out_tokens, _ = generate([tokens], model, max_tokens=512, temperature=0.0, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
-    summarized_text = tokenizer.instruct_tokenizer.tokenizer.decode(out_tokens[0])
-    return summarized_text
+# Crear la cadena LLM
+llm_chain = LLMChain(llm=mistral_llm, prompt=prompt)
 
+# Función para manejar archivos subidos
 def handle_uploaded_file(uploaded_file):
     try:
         if uploaded_file.name.endswith(".txt"):
@@ -104,66 +81,109 @@ def handle_uploaded_file(uploaded_file):
     except Exception as e:
         return str(e)
 
+# Función para traducir texto
+def translate(text, target_language):
+    context = ""
+    question = f"Por favor, traduzca el siguiente documento al {target_language}:\n{text}\nAsegúrese de que la traducción sea precisa y conserve el significado original del documento."
+    response = llm_chain.run(context=context, question=question)
+    return response
+
+# Función para resumir texto
+def summarize(text, length):
+    context = ""
+    question = f"Por favor, haga un resumen {length} del siguiente documento:\n{text}\nAsegúrese de que el resumen sea conciso y conserve el significado original del documento."
+    response = llm_chain.run(context=context, question=question)
+    return response
+
+# Configuración del modelo de clasificación
+@st.cache_resource
+def load_classification_model():
+    tokenizer_cls = AutoTokenizer.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
+    model_cls = AutoModelForSequenceClassification.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
+    return model_cls, tokenizer_cls
+
+classification_model, classification_tokenizer = load_classification_model()
+
+id2label = {0: "multas", 1: "politicas_de_privacidad", 2: "contratos", 3: "denuncias", 4: "otros"}
+
+def classify_text(text):
+    inputs = classification_tokenizer(text, return_tensors="pt", max_length=4096, truncation=True, padding="max_length")
+    classification_model.eval()
+    with torch.no_grad():
+        outputs = classification_model(**inputs)
+    logits = outputs.logits
+    predicted_class_id = logits.argmax(dim=-1).item()
+    predicted_label = id2label[predicted_class_id]
+    return predicted_label
+
+# Función para cargar documentos JSON
+def load_json_documents(category):
+    try:
+        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
+    except FileNotFoundError:
+        return []
+
+# Configuración de FAISS y embeddings
+@st.cache_resource
+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 explain_text(user_input, document_text):
+    classification = classify_text(document_text)
+    if classification in ["multas", "politicas_de_privacidad", "contratos", "denuncias"]:
+        docs = load_json_documents(classification)
+        if docs:
+            vector_store = create_vector_store(docs)
+            search_docs = vector_store.similarity_search(user_input)
+            context = " ".join([doc.page_content for doc in search_docs])
+        else:
+            context = ""
+    else:
+        context = ""
+    question = user_input
+    response = llm_chain.run(context=context, question=question)
+    return response
+
 def main():
     st.title("LexAIcon")
     st.write("Puedes conversar con este chatbot basado en Mistral-7B-Instruct y subir archivos para que el chatbot los procese.")
 
-    if "messages" not in st.session_state:
-        st.session_state["messages"] = []
-
     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)")
 
-    user_input = st.text_input("Introduce tu consulta:", "")
-    
-    if user_input:
-        st.session_state.messages.append({"role": "user", "content": user_input})
-
-        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 == "Resumir":
-            summary_length = st.selectbox("Selecciona la longitud del resumen", ["corto", "medio", "largo"])
-
-        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)
-                classification = classify_text(file_content)
-                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])
-                completion_request = ChatCompletionRequest(
-                    messages=[UserMessage(content=f"Contexto: {context}\n\nPregunta: {user_input}")]
-                )
-                tokens = tokenizer.encode_chat_completion(completion_request).tokens
-                out_tokens, _ = generate([tokens], model, max_tokens=512, temperature=0.0, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
-                bot_response = tokenizer.instruct_tokenizer.tokenizer.decode(out_tokens[0])
-        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"
-            bot_response = summarize(user_input, length)
-        elif operation == "Traducir":
-            bot_response = translate(user_input, target_language)
-        else:
-            completion_request = ChatCompletionRequest(
-                messages=[UserMessage(content=user_input)]
-            )
-            tokens = tokenizer.encode_chat_completion(completion_request).tokens
-            out_tokens, _ = generate([tokens], model, max_tokens=512, temperature=0.0, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
-            bot_response = tokenizer.instruct_tokenizer.tokenizer.decode(out_tokens[0])
-
-        st.session_state.messages.append({"role": "assistant", "content": bot_response})
-        st.write(f"**Assistant:** {bot_response}")
+    operation = st.radio("Selecciona una operación", ["Resumir", "Traducir", "Explicar"])
+
+    if operation == "Explicar":
+        user_input = st.text_area("Introduce tu pregunta:", "")
+        uploaded_file = st.file_uploader("Sube un archivo", type=["txt", "pdf", "docx", "csv", "json"])
+        if uploaded_file and user_input:
+            document_text = handle_uploaded_file(uploaded_file)
+            bot_response = explain_text(user_input, document_text)
+            st.write(f"**Assistant:** {bot_response}")
+    else:
+        uploaded_file = st.file_uploader("Sube un archivo", type=["txt", "pdf", "docx", "csv", "json"])
+        if uploaded_file:
+            document_text = handle_uploaded_file(uploaded_file)
+            if operation == "Traducir":
+                target_language = st.selectbox("Selecciona el idioma de traducción", ["español", "inglés", "francés", "alemán"])
+                bot_response = translate(document_text, target_language)
+            elif operation == "Resumir":
+                summary_length = st.selectbox("Selecciona la longitud del resumen", ["corto", "medio", "largo"])
+                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"
+                bot_response = summarize(document_text, length)
+            st.write(f"**Assistant:** {bot_response}")
 
 if __name__ == "__main__":
     main()