Update app.py

app.py

@@ -1,17 +1,15 @@
 import gradio as gr
 import os
 
-from langchain_community.document_loaders import PyPDFLoader
+from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import Chroma
+from langchain.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
-from langchain_community.embeddings import HuggingFaceEmbeddings 
-from langchain_community.llms import HuggingFacePipeline
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.llms import HuggingFacePipeline
 from langchain.memory import ConversationBufferMemory
 
 from transformers import AutoTokenizer, pipeline
-import torch
-
 import chromadb
 
 # Lista de modelos 100% abertos e gratuitos
@@ -32,12 +30,12 @@ list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
 # Função para carregar documentos PDF
 def load_doc(list_file_path, chunk_size, chunk_overlap):
-    loaders = [PyPDFLoader(x) for x in list_file_path]
+    loaders = [PyPDFLoader(file_path) for file_path in list_file_path]
     pages = []
     for loader in loaders:
         pages.extend(loader.load())
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=chunk_size, 
+        chunk_size=chunk_size,
         chunk_overlap=chunk_overlap
     )
     return text_splitter.split_documents(pages)
@@ -48,8 +46,7 @@ def create_db(splits, collection_name):
     return Chroma.from_documents(
         documents=splits,
         embedding=embedding,
-        client=chromadb.PersistentClient(),
-        collection_name=collection_name
+        persist_directory=f"./{collection_name}"
     )
 
 # Função para inicializar o modelo LLM
@@ -63,8 +60,8 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
         "text-generation",
         model=llm_model,
         tokenizer=tokenizer,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        device=0 if torch.cuda.is_available() else -1,
         max_new_tokens=max_tokens,
         do_sample=True,
         top_k=top_k,
@@ -90,8 +87,8 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
 # Interface Gradio
 def demo():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        vector_db = gr.State()
-        qa_chain = gr.State()
+        vector_db = gr.State(None)  # Inicializa com None
+        qa_chain = gr.State(None)  # Inicializa com None
         
         gr.Markdown("## 🤖 Chatbot para PDFs com Modelos Gratuitos")
         
@@ -116,14 +113,26 @@ def demo():
             clear_btn = gr.Button("Limpar Chat")
             
         # Eventos
+        def process_documents(files, cs, co):
+            file_paths = [f.name for f in files]
+            splits = load_doc(file_paths, cs, co)
+            db = create_db(splits, "docs")
+            return db, "Documentos processados!"
+        
         process_btn.click(
-            lambda files, cs, co: (create_db(load_doc([f.name for f in files], cs, co), "docs"), "Documentos processados!"),
+            process_documents,
             inputs=[pdf_input, chunk_size, chunk_overlap],
             outputs=[vector_db, process_status]
         )
         
+        def load_model(model, temp):
+            if vector_db.value is None:
+                return None, "Por favor, processe os documentos primeiro."
+            qa = initialize_llmchain(list_llm[list_llm_simple.index(model)], temp, 512, 3, vector_db.value)
+            return qa, "Modelo carregado!"
+        
         load_model_btn.click(
-            lambda model, temp: (initialize_llmchain(list_llm[list_llm_simple.index(model)], temp, 512, 3, vector_db.value), "Modelo carregado!"),
+            load_model,
             inputs=[model_selector, temperature],
             outputs=[qa_chain, model_status]
         )
@@ -147,4 +156,4 @@ def demo():
     demo.launch()
 
 if __name__ == "__main__":
-    demo()
+    demo()