Update app.py

app.py

@@ -9,27 +9,23 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFacePipeline
 from langchain.memory import ConversationBufferMemory
 
-from pathlib import Path
-import chromadb
-from unidecode import unidecode
-
 from transformers import AutoTokenizer, pipeline
-import transformers
 import torch
-import re
+
+import chromadb
 
 # Lista de modelos 100% abertos e gratuitos
 list_llm = [
-    "google/flan-t5-xxl",          # Modelo para tarefas text-to-text
-    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",  # Modelo leve para diálogo
-    "microsoft/phi-2",             # Modelo para raciocínio lógico
-    "facebook/opt-1.3b",           # Modelo de geração de texto
-    "EleutherAI/gpt-neo-1.3B",     # Versão open-source do GPT-3
-    "bigscience/bloom-1b7",        # Modelo multilíngue
-    "RWKV/rwkv-4-169m-pile",       # Modelo eficiente em RAM
-    "gpt2-medium",                 # Clássico modelo de GPT-2
-    "databricks/dolly-v2-3b",      # Modelo para instruções
-    "mosaicml/mpt-7b-instruct"     # Modelo para instruções
+    "google/flan-t5-xxl",
+    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    "microsoft/phi-2",
+    "facebook/opt-1.3b",
+    "EleutherAI/gpt-neo-1.3B",
+    "bigscience/bloom-1b7",
+    "RWKV/rwkv-4-169m-pile",
+    "gpt2-medium",
+    "databricks/dolly-v2-3b",
+    "mosaicml/mpt-7b-instruct"
 ]
 
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
@@ -52,7 +48,7 @@ def create_db(splits, collection_name):
     return Chroma.from_documents(
         documents=splits,
         embedding=embedding,
-        client=chromadb.EphemeralClient(),
+        client=chromadb.PersistentClient(),
         collection_name=collection_name
     )
 
@@ -88,7 +84,6 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
         llm=llm,
         retriever=vector_db.as_retriever(),
         memory=memory,
-        chain_type="stuff",
         return_source_documents=True
     )
 
@@ -118,7 +113,7 @@ def demo():
         with gr.Tab("💬 Chat"):
             chatbot = gr.Chatbot(height=400)
             msg = gr.Textbox(label="Sua mensagem")
-            clear_btn = gr.ClearButton([msg, chatbot])
+            clear_btn = gr.Button("Limpar Chat")
             
         # Eventos
         process_btn.click(
@@ -136,16 +131,20 @@ def demo():
         def respond(message, chat_history):
             if qa_chain.value is None:
                 return "Por favor, carregue um modelo primeiro.", chat_history
+            
             result = qa_chain.value({"question": message, "chat_history": chat_history})
             response = result["answer"]
+            
             sources = "\n".join([f"📄 Página {doc.metadata['page']+1}: {doc.page_content[:50]}..." 
-                               for doc in result["source_documents"][:2]])
+                                   for doc in result.get("source_documents", [])[:2]])
+            
             chat_history.append((message, f"{response}\n\n🔍 Fontes:\n{sources}"))
             return "", chat_history
         
         msg.submit(respond, [msg, chatbot], [msg, chatbot])
+        clear_btn.click(lambda: [], outputs=[chatbot])
         
     demo.launch()
 
 if __name__ == "__main__":
-    demo()
+    demo()