Update app.py

app.py

@@ -1,4 +1,11 @@
 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 PyPDF2 import PdfReader
 from docx import Document
@@ -9,28 +16,22 @@ import torch
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from huggingface_hub import login, InferenceClient
 
-huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
+# 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)
 
-# Realizar el inicio de sesión de Hugging Face solo si el token está disponible
-if huggingface_token:
-    login(token=huggingface_token)
-
-# Configuración del cliente de inferencia
-@st.cache_resource
-def load_inference_client():
-    client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.3")
-    return client
-
-client = load_inference_client()
+# Configurar el modelo y el tokenizador
+tokenizer = MistralTokenizer.from_file(f"{mistral_models_path}/tokenizer.model.v3")
+model = Transformer.from_folder(mistral_models_path)
 
 # Configuración del modelo de clasificación
 @st.cache_resource
 def load_classification_model():
-    tokenizer = AutoTokenizer.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
-    model = AutoModelForSequenceClassification.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
-    return model, tokenizer
+    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()
 
@@ -60,29 +61,21 @@ def create_vector_store(docs):
     return vector_store
 
 def translate(text, target_language):
-    template = f'''
-    Por favor, traduzca el siguiente documento al {target_language}:
-<document>
-{text}
-</document>
-Asegúrese de que la traducción sea precisa y conserve el significado original del documento.
-    '''
-    messages = [{"role": "user", "content": template}]
-    response = client.chat(messages)
-    translated_text = response.generated_text
+    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
 
 def summarize(text, length):
-    template = f'''
-    Por favor, haga un resumen {length} del siguiente documento:
-<document>
-{text}
-</document>
-Asegúrese de que el resumen sea conciso y conserve el significado original del documento.
-    '''
-    messages = [{"role": "user", "content": template}]
-    response = client.chat(messages)
-    summarized_text = response.generated_text
+    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
 
 def handle_uploaded_file(uploaded_file):
@@ -122,9 +115,6 @@ def main():
         st.text_input("HuggingFace Token", value=huggingface_token, type="password", key="huggingface_token")
         st.caption("[Consigue un HuggingFace Token](https://huggingface.co/settings/tokens)")
 
-    for msg in st.session_state.messages:
-        st.write(f"**{msg['role'].capitalize()}:** {msg['content']}")
-
     user_input = st.text_input("Introduce tu consulta:", "")
     
     if user_input:
@@ -148,10 +138,12 @@ def main():
                 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}"
-                messages = [{"role": "user", "content": prompt_with_context}]
-                response = client.chat(messages)
-                bot_response = response.generated_text
+                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"
@@ -163,9 +155,12 @@ def main():
         elif operation == "Traducir":
             bot_response = translate(user_input, target_language)
         else:
-            messages = [{"role": "user", "content": user_input}]
-            response = client.chat(messages)
-            bot_response = response.generated_text
+            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}")