Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
+from transformers import InferenceClient
 from PyPDF2 import PdfReader
 from docx import Document
 import csv
@@ -10,6 +10,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from huggingface_hub import login
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
 
@@ -17,15 +18,13 @@ huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
 if huggingface_token:
     login(token=huggingface_token)
 
-# Cargar el modelo y el pipeline de Hugging Face
+# Configuración del cliente de inferencia
 @st.cache_resource
-def load_pipeline():
-    tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
-    model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
-    text_gen_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    return text_gen_pipeline
+def load_inference_client():
+    client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.3")
+    return client
 
-text_gen_pipeline = load_pipeline()
+client = load_inference_client()
 
 # Configuración del modelo de clasificación
 @st.cache_resource
@@ -75,18 +74,16 @@ def classify_text(text):
     return predicted_label
 
 def translate(text, target_language):
-    template = '''
-    Por favor, traduzca el siguiente documento al {LANGUAGE}:
+    template = f'''
+    Por favor, traduzca el siguiente documento al {target_language}:
 <document>
-{TEXT}
+{text}
 </document>
 Asegúrese de que la traducción sea precisa y conserve el significado original del documento.
     '''
-    
-    formatted_prompt = template.replace("{TEXT}", text).replace("{LANGUAGE}", target_language)
-    response = text_gen_pipeline(formatted_prompt, max_length=512)
-    translated_text = response[0]['generated_text']
-    
+    messages = [{"role": "user", "content": template}]
+    response = client.chat(messages)
+    translated_text = response.generated_text
     return translated_text
 
 def summarize(text, length):
@@ -97,10 +94,9 @@ def summarize(text, length):
 </document>
 Asegúrese de que el resumen sea conciso y conserve el significado original del documento.
     '''
-    
-    response = text_gen_pipeline(template, max_length=512)
-    summarized_text = response[0]['generated_text']
-    
+    messages = [{"role": "user", "content": template}]
+    response = client.chat(messages)
+    summarized_text = response.generated_text
     return summarized_text
 
 def handle_uploaded_file(uploaded_file):
@@ -131,21 +127,22 @@ def handle_uploaded_file(uploaded_file):
 
 def main():
     st.title("LexAIcon")
-    st.write("Puedes conversar con este chatbot basado en Mistral7B-Instruct y subir archivos para que el chatbot los procese.")
+    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"] = [{"role": "assistant", "content": "¿Cómo puedo ayudarte?"}]
+        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)")
 
     for msg in st.session_state.messages:
-        st.chat_message(msg["role"]).write(msg["content"])
+        st.write(f"**{msg['role'].capitalize()}:** {msg['content']}")
 
-    if prompt := st.chat_input():
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        st.chat_message("user").write(prompt)
+    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
@@ -162,12 +159,12 @@ def main():
                 file_content = handle_uploaded_file(uploaded_file)
                 classification = classify_text(file_content)
                 vector_store = vector_stores[classification]
-                search_docs = vector_store.similarity_search(prompt)
+                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: {prompt}"
-                response = text_gen_pipeline(prompt_with_context, max_length=512)
-                msg = response[0]['generated_text']
-
+                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
         elif operation == "Resumir":
             if summary_length == "corto":
                 length = "de aproximadamente 50 palabras"
@@ -175,17 +172,16 @@ def main():
                 length = "de aproximadamente 100 palabras"
             elif summary_length == "largo":
                 length = "de aproximadamente 500 palabras"
-            msg = summarize(prompt, length)
-
+            bot_response = summarize(user_input, length)
         elif operation == "Traducir":
-            msg = translate(prompt, target_language)
-
+            bot_response = translate(user_input, target_language)
         else:
-            response = text_gen_pipeline(prompt, max_length=512)
-            msg = response[0]['generated_text']
+            messages = [{"role": "user", "content": user_input}]
+            response = client.chat(messages)
+            bot_response = response.generated_text
 
-        st.session_state.messages.append({"role": "assistant", "content": msg})
-        st.chat_message("assistant").write(msg)
+        st.session_state.messages.append({"role": "assistant", "content": bot_response})
+        st.write(f"**Assistant:** {bot_response}")
 
 if __name__ == "__main__":
     main()