Update app.py

app.py

@@ -3,67 +3,84 @@ from transformers import pipeline
 from langchain_community.llms import HuggingFacePipeline
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain_community.document_loaders import PyPDFLoader
+from langchain.document_loaders import TextLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chains import ConversationalRetrievalChain
-from langchain.memory import ConversationBufferMemory
+from langchain.prompts import PromptTemplate
 
-# Carrega e divide o documento PDF fixo (local)
-def carregar_documento_pdf(path):
-    loader = PyPDFLoader(path)
-    documentos = loader.load()
-    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
-    documentos_divididos = splitter.split_documents(documentos)
-    embeddings = HuggingFaceEmbeddings()
-    vectorstore = FAISS.from_documents(documentos_divididos, embeddings)
-    return vectorstore
-
-# Inicializa vectorstore com PDF fixo
-vectorstore = carregar_documento_pdf("meu_documento.pdf")
-
-# Inicializa pipeline do Flan-Alpaca
+# Configura o pipeline do modelo
 pipe = pipeline(
     "text2text-generation",
     model="declare-lab/flan-alpaca-base",
-    max_new_tokens=100,
+    max_new_tokens=150,
     truncation=True,
     max_length=512,
 )
 llm = HuggingFacePipeline(pipeline=pipe)
 
-# Memória limitada para evitar prompts muito longos (k=3 últimas mensagens)
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True, k=3)
+# Carrega e divide documento em chunks maiores para contexto melhor
+def carregar_documento_txt(path):
+    loader = TextLoader(path, encoding="latin-1")
+    documentos = loader.load()
+    # Chunk size maior e overlap para contexto mais coeso
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    documentos_divididos = splitter.split_documents(documentos)
+    embeddings = HuggingFaceEmbeddings()
+    vectorstore = FAISS.from_documents(documentos_divididos, embeddings)
+    return vectorstore
+
+vectorstore = carregar_documento_txt("meu_documento.txt")
 
-# Cadeia de Perguntas & Respostas com memória
+# Prompt customizado para garantir que o modelo responda com base no contexto
+template = """
+Você é um assistente que responde perguntas baseado no contexto a seguir.
+Se não souber a resposta, responda "Não sei com base nos documentos fornecidos."
+Contexto:
+{context}
+
+Pergunta:
+{question}
+
+Resposta:
+"""
+QA_CHAIN_PROMPT = PromptTemplate(template=template, input_variables=["context", "question"])
+
+# Cria o chain de busca sem memória
 chat_chain = ConversationalRetrievalChain.from_llm(
     llm=llm,
     retriever=vectorstore.as_retriever(),
-    memory=memory,
+    memory=None,  # sem memória para simplificar
+    combine_docs_chain_kwargs={"prompt": QA_CHAIN_PROMPT}
 )
 
-# Função para responder pergunta com contexto de memória
-def responder_pergunta(pergunta):
-    if not pergunta or not pergunta.strip():
-        return "Por favor, insira uma pergunta válida."
+# Função para responder perguntas com base no documento
+def responder_documento(pergunta):
     resposta = chat_chain.run({"question": pergunta})
+    # Pós-processar: se resposta muito curta ou igual ao contexto, retorna fallback
+    if len(resposta.strip()) < 10 or resposta.lower().startswith("não sei"):
+        return "Não sei com base nos documentos fornecidos."
     return resposta
 
-# Interface Gradio com chat e memória
+# Interface Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## 🤖 Chat com PDF fixo - Flan Alpaca + LangChain + FAISS + Memória")
-    chatbot = gr.Chatbot(label="Conversa com memória")
-    pergunta_input = gr.Textbox(placeholder="Digite sua pergunta aqui...", label="Pergunta")
-    enviar_btn = gr.Button("Enviar")
+    gr.Markdown("## 🤖 Chat IA com Documento TXT - Flan-Alpaca e LangChain")
+
+    chatbot = gr.Chatbot(label="Chat com Documento")
+    entrada = gr.Textbox(placeholder="Digite sua pergunta e aperte Enter", show_label=False)
+    estado_chat = gr.State([])
 
-    def interagir(pergunta, chat_hist):
-        resposta = responder_pergunta(pergunta)
-        chat_hist = chat_hist + [(pergunta, resposta)]
-        return "", chat_hist
+    def interagir(msg, history):
+        try:
+            resposta = responder_documento(msg)
+        except Exception as e:
+            resposta = f"Erro ao gerar resposta: {e}"
+        history = history + [(msg, resposta)]
+        return "", history
 
-    enviar_btn.click(interagir, inputs=[pergunta_input, chatbot], outputs=[pergunta_input, chatbot])
-    pergunta_input.submit(interagir, inputs=[pergunta_input, chatbot], outputs=[pergunta_input, chatbot])
+    entrada.submit(interagir, [entrada, estado_chat], [entrada, chatbot])
+    gr.Button("Enviar").click(interagir, [entrada, estado_chat], [entrada, chatbot])
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)