Update app.py

app.py

@@ -13,24 +13,15 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 
-huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
-
-# 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 modelo de generación de texto
+# Cargar el modelo y el pipeline de Hugging Face
 @st.cache_resource
-def load_llm():
-    llm = HuggingFaceEndpoint(
-        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-        task="text-generation"
-    )
-    llm_engine_hf = ChatHuggingFace(llm=llm)
+def load_pipeline():
     tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
-    return llm_engine_hf, tokenizer
+    model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+    text_gen_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    return text_gen_pipeline
 
-llm_engine_hf, tokenizer = load_llm()
+text_gen_pipeline = load_pipeline()
 
 # Configuración del modelo de clasificación
 @st.cache_resource
@@ -89,8 +80,8 @@ Asegúrese de que la traducción sea precisa y conserve el significado original
     '''
     
     formatted_prompt = template.replace("{TEXT}", text).replace("{LANGUAGE}", target_language)
-    response = llm_engine_hf.invoke(formatted_prompt)
-    translated_text = response.content
+    response = text_gen_pipeline(formatted_prompt, max_length=512)
+    translated_text = response[0]['generated_text']
     
     return translated_text
 
@@ -103,8 +94,8 @@ def summarize(text, length):
 Asegúrese de que el resumen sea conciso y conserve el significado original del documento.
     '''
     
-    response = llm_engine_hf.invoke(template)
-    summarized_text = response.content
+    response = text_gen_pipeline(template, max_length=512)
+    summarized_text = response[0]['generated_text']
     
     return summarized_text
 
@@ -135,7 +126,6 @@ def handle_uploaded_file(uploaded_file):
         return str(e)
 
 def main():
-    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.")
 
@@ -171,8 +161,8 @@ def main():
                 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
+                response = text_gen_pipeline(prompt_with_context, max_length=512)
+                msg = response[0]['generated_text']
 
         elif operation == "Resumir":
             if summary_length == "corto":
@@ -187,7 +177,8 @@ def main():
             msg = translate(prompt, target_language)
 
         else:
-            msg = llm_engine_hf.invoke(prompt).content
+            response = text_gen_pipeline(prompt, max_length=512)
+            msg = response[0]['generated_text']
 
         st.session_state.messages.append({"role": "assistant", "content": msg})
         st.chat_message("assistant").write(msg)