Update app.py

app.py

@@ -13,7 +13,64 @@ from langchain_community.llms import HuggingFaceEndpoint
 list_llm = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.2"]  
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
-# Funções existentes (load_doc, create_db, initialize_llmchain, etc.) permanecem iguais...
+# Load and split PDF document
+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)
+    doc_splits = text_splitter.split_documents(pages)
+    return doc_splits
+
+# Create vector database
+def create_db(splits):
+    embeddings = HuggingFaceEmbeddings()
+    vectordb = FAISS.from_documents(splits, embeddings)
+    return vectordb
+
+# Initialize database
+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 langchain LLM chain
+def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
+    if llm_model == "meta-llama/Meta-Llama-3-8B-Instruct":
+        llm = HuggingFaceEndpoint(
+            repo_id=llm_model,
+            huggingfacehub_api_token=api_token,
+            temperature=temperature,
+            max_new_tokens=max_tokens,
+            top_k=top_k,
+        )
+    else:
+        llm = HuggingFaceEndpoint(
+            huggingfacehub_api_token=api_token,
+            repo_id=llm_model, 
+            temperature=temperature,
+            max_new_tokens=max_tokens,
+            top_k=top_k,
+        )
+    memory = ConversationBufferMemory(memory_key="chat_history", output_key='answer', return_messages=True)
+    retriever = vector_db.as_retriever()
+    qa_chain = ConversationalRetrievalChain.from_llm(
+        llm,
+        retriever=retriever,
+        chain_type="stuff", 
+        memory=memory,
+        return_source_documents=True,
+        verbose=False,
+    )
+    return qa_chain
+
+# Initialize LLM
+def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
+    llm_name = list_llm[llm_option]
+    qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress)
+    return qa_chain, "QA chain initialized. Chatbot is ready!"
 
 def format_chat_history(message, chat_history):
     formatted_chat_history = []
@@ -22,10 +79,8 @@ def format_chat_history(message, chat_history):
         formatted_chat_history.append(f"Assistant: {bot_message}")
     return formatted_chat_history
 
-# Ajuste na função conversation para suportar idioma
 def conversation(qa_chain, message, history, language):
     formatted_chat_history = format_chat_history(message, history)
-    # Generate response using QA chain
     response = qa_chain.invoke({"question": message, "chat_history": formatted_chat_history})
     response_answer = response["answer"]
     if response_answer.find("Helpful Answer:") != -1:
@@ -33,8 +88,6 @@ def conversation(qa_chain, message, history, language):
     
     # Ajustar resposta com base no idioma
     if language == "Português":
-        # Aqui, idealmente, você usaria uma API de tradução ou o modelo geraria diretamente em português
-        # Como exemplo, adiciono uma mensagem fixa para demonstrar
         response_answer = f"Resposta em português: {response_answer}"
     else:
         response_answer = f"Response in English: {response_answer}"
@@ -74,7 +127,6 @@ def demo():
 
             with gr.Column(scale=200):
                 gr.Markdown("<b>Step 2 - Chat with your Document</b>")
-                # Adicionar seletor de idioma
                 language_selector = gr.Radio(["English", "Português"], label="Select Language", value="English")
                 chatbot = gr.Chatbot(height=505)
                 with gr.Accordion("Relevant context from the source document", open=False):
@@ -94,7 +146,7 @@ def demo():
             lambda: [None, "", 0, "", 0, "", 0], inputs=None, outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], queue=False
         )
 
-        # Chatbot events com o idioma
+        # Chatbot events
         msg.submit(conversation, inputs=[qa_chain, msg, chatbot, language_selector], 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, language_selector], 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)