Update app.py

app.py

@@ -11,41 +11,40 @@ from langchain.memory import ConversationBufferMemory
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# List of local models (publicly accessible, no token required)
-list_llm = ["facebook/opt-350m", "distilbert/distilgpt2"]  
+# Lista de modelos públicos e leves
+list_llm = ["EleutherAI/gpt-neo-125m", "distilbert/distilgpt2"]  
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
-# Load and split PDF document
+# Carregar e dividir documento PDF
 def load_doc(list_file_path):
     loaders = [PyPDFLoader(x) for x in list_file_path]
     pages = []
     for loader in loaders:
         pages.extend(loader.load())
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1024, 
-        chunk_overlap=64 
+        chunk_size=512,  # Reduzido para acelerar a busca
+        chunk_overlap=32  # Menor sobreposição para menos processamento
     )  
     doc_splits = text_splitter.split_documents(pages)
     return doc_splits
 
-# Create vector database
+# Criar banco de vetores
 def create_db(splits):
     embeddings = HuggingFaceEmbeddings()
     vectordb = FAISS.from_documents(splits, embeddings)
     return vectordb
 
-# Initialize langchain LLM chain with local model
+# Inicializar o chain LLM local
 def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
-    # Load the model and tokenizer locally
     tokenizer = AutoTokenizer.from_pretrained(llm_model)
     model = AutoModelForCausalLM.from_pretrained(
         llm_model,
-        device_map="auto",  # Automatically use GPU if available
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Optimize for GPU or CPU
-        trust_remote_code=True  # Required for some models
+        device_map="auto",  # Usa GPU se disponível
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Otimiza para GPU
+        trust_remote_code=True
     )
 
-    # Create a pipeline for text generation
+    # Pipeline otimizado
     pipe = pipeline(
         "text-generation",
         model=model,
@@ -53,12 +52,11 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
         max_new_tokens=max_tokens,
         temperature=temperature,
         top_k=top_k,
-        do_sample=True,
+        do_sample=False,  # Greedy decoding para mais velocidade
         repetition_penalty=1.1,
         return_full_text=False
     )
 
-    # Wrap the pipeline in HuggingFacePipeline for LangChain
     llm = HuggingFacePipeline(pipeline=pipe)
     
     memory = ConversationBufferMemory(
@@ -67,7 +65,7 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
         return_messages=True
     )
 
-    retriever = vector_db.as_retriever()
+    retriever = vector_db.as_retriever(search_kwargs={"k": 2})  # Reduzir número de documentos retornados
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
         retriever=retriever,
@@ -78,14 +76,14 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
     )
     return qa_chain
 
-# Initialize database
+# Inicializar banco de dados
 def initialize_database(list_file_obj, progress=gr.Progress()):
     list_file_path = [x.name for x in list_file_obj if x is not None]
     doc_splits = load_doc(list_file_path)
     vector_db = create_db(doc_splits)
     return vector_db, "Database created!"
 
-# Initialize LLM
+# Inicializar LLM
 def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
     llm_name = list_llm[llm_option]
     print("llm_name: ", llm_name)
@@ -108,10 +106,10 @@ def conversation(qa_chain, message, history):
     response_sources = response["source_documents"]
     response_source1 = response_sources[0].page_content.strip()
     response_source2 = response_sources[1].page_content.strip()
-    response_source3 = response_sources[2].page_content.strip()
+    response_source3 = ""  # Menos referências para acelerar
     response_source1_page = response_sources[0].metadata["page"] + 1
     response_source2_page = response_sources[1].metadata["page"] + 1
-    response_source3_page = response_sources[2].metadata["page"] + 1
+    response_source3_page = 0
     new_history = history + [(message, response_answer)]
     return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
 
@@ -127,7 +125,7 @@ def demo():
         vector_db = gr.State()
         qa_chain = gr.State()
         gr.HTML("<center><h1>RAG PDF chatbot</h1><center>")
-        gr.Markdown("""<b>Query your PDF documents!</b> This AI agent is designed to perform retrieval augmented generation (RAG) on PDF documents. This version runs locally and does not require an API token. \
+        gr.Markdown("""<b>Query your PDF documents!</b> This AI agent is designed to perform retrieval augmented generation (RAG) on PDF documents. Optimized for speed without an API token. \
         <b>Please do not upload confidential documents.</b>
         """)
         with gr.Row():
@@ -145,11 +143,11 @@ def demo():
                 with gr.Row():
                     with gr.Accordion("LLM input parameters", open=False):
                         with gr.Row():
-                            slider_temperature = gr.Slider(minimum=0.01, maximum=1.0, value=0.5, step=0.1, label="Temperature", info="Controls randomness in token generation", interactive=True)
+                            slider_temperature = gr.Slider(minimum=0.01, maximum=1.0, value=0.5, step=0.1, label="Temperature", info="Controls randomness (ignored with greedy decoding)", interactive=True)
                         with gr.Row():
-                            slider_maxtokens = gr.Slider(minimum=128, maximum=2048, value=512, step=128, label="Max New Tokens", info="Maximum number of tokens to be generated", interactive=True)
+                            slider_maxtokens = gr.Slider(minimum=64, maximum=512, value=128, step=64, label="Max New Tokens", info="Maximum number of tokens to be generated", interactive=True)
                         with gr.Row():
-                            slider_topk = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="top-k", info="Number of tokens to select the next token from", interactive=True)
+                            slider_topk = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="top-k", info="Number of tokens to select (ignored with greedy decoding)", interactive=True)
                 with gr.Row():
                     qachain_btn = gr.Button("Initialize Question Answering Chatbot")
                 with gr.Row():
@@ -174,7 +172,7 @@ def demo():
                     submit_btn = gr.Button("Submit")
                     clear_btn = gr.ClearButton([msg, chatbot], value="Clear")
             
-        # Preprocessing events
+        # Eventos de pré-processamento
         db_btn.click(initialize_database, inputs=[document], outputs=[vector_db, db_progress])
         qachain_btn.click(initialize_LLM, inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], outputs=[qa_chain, llm_progress]).then(
             lambda: [None, "", 0, "", 0, "", 0],
@@ -183,7 +181,7 @@ def demo():
             queue=False
         )
 
-        # Chatbot events
+        # Eventos do chatbot
         msg.submit(conversation, inputs=[qa_chain, msg, chatbot], outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], queue=False)
         submit_btn.click(conversation, inputs=[qa_chain, msg, chatbot], outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], queue=False)
         clear_btn.click(lambda: [None, "", 0, "", 0, "", 0], inputs=None, outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], queue=False)