Update app.py

app.py

@@ -1,20 +1,15 @@
 import gradio as gr
 import os
 
-from langchain.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import Chroma
-from langchain.chains import ChatVectorDBChain
-from langchain.embeddings import OpenAIEmbeddings
+from langchain_community.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
 from langchain_community.embeddings import HuggingFaceEmbeddings 
 from langchain_community.llms import HuggingFacePipeline
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
 from langchain_community.llms import HuggingFaceEndpoint
-from langchain_community.chat_models.openai import ChatOpenAI
-from langchain.chains import ConversationalRetrievalChain
-from langchain.llms import OpenAI
 
 from pathlib import Path
 import chromadb
@@ -26,12 +21,27 @@ import torch
 import tqdm 
 import accelerate
 import re
+from openai import OpenAI
+from huggingface_hub import HfApi
+from huggingface_hub import hf_hub_download
+
+from langchain.document_loaders import PyPDFLoader
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.chains import ChatVectorDBChain
+from langchain.vectorstores import Chroma
+from langchain.llms import OpenAI
+import os
 
 
+os.environ["OPENAI_API_KEY"]
 TOKEN = os.getenv('HUGGING_FACE_HUB_TOKEN')
 
 
-    
+
+# default_persist_directory = './chroma_HF/'
+list_llm = ["jondurbin/airoboros-gpt-3.5-turbo-100k-7b"]
+list_llm_simple = [os.path.basename(llm) for llm in list_llm]
+
 # Load PDF document and create doc splits
 def load_doc(list_file_path, chunk_size, chunk_overlap):
     # Processing for one document only
@@ -49,17 +59,20 @@ def load_doc(list_file_path, chunk_size, chunk_overlap):
     return doc_splits
 
 
-# Initialize your vector database with OpenAIEmbeddings and persist it
-def create_db(documents, collection_name, persist_directory="."):
-    embeddings = OpenAIEmbeddings()
+# Create vector database
+def create_db(splits, collection_name):
+    embedding = OpenAIEmbeddings()
+    new_client = chromadb.EphemeralClient()
     vectordb = Chroma.from_documents(
-        documents=documents,
-        embedding=embeddings,
+        documents=splits,
+        embedding=embedding,
+        persist_directory=".",
+        client=new_client,
         collection_name=collection_name,
-        persist_directory=persist_directory
+        # persist_directory=default_persist_directory
     )
-    vectordb.persist()
-    return vectordb
+    return vectordb.persist()
+
 
 # Load vector database
 def load_db():
@@ -70,34 +83,74 @@ def load_db():
     return vectordb
 
 
+# Initialize langchain LLM chain
+def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
+    progress(0.1, desc="Initializing HF tokenizer...")
+    # HuggingFacePipeline uses local model
+    # Note: it will download model locally...
+    # tokenizer=AutoTokenizer.from_pretrained(llm_model)
+    # progress(0.5, desc="Initializing HF pipeline...")
+    # pipeline=transformers.pipeline(
+    #     "text-generation",
+    #     model=llm_model,
+    #     tokenizer=tokenizer,
+    #     torch_dtype=torch.bfloat16,
+    #     trust_remote_code=True,
+    #     device_map="auto",
+    #     # max_length=1024,
+    #     max_new_tokens=max_tokens,
+    #     do_sample=True,
+    #     top_k=top_k,
+    #     num_return_sequences=1,
+    #     eos_token_id=tokenizer.eos_token_id
+    #     )
+    # llm = HuggingFacePipeline(pipeline=pipeline, model_kwargs={'temperature': temperature})
+    
+    # HuggingFaceHub uses HF inference endpoints
+    progress(0.5, desc="Initializing HF Hub...")
+    if llm_model == "jondurbin/airoboros-gpt-3.5-turbo-100k-7b":
+        llm = HuggingFaceEndpoint(
+            repo_id=llm_model, 
+            # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "load_in_8bit": True}
+            temperature = temperature,
+            max_new_tokens = max_tokens,
+            top_k = top_k,
+            load_in_8bit = True,
+        )
 
-
-# Assuming vectordb is correctly initialized and persisted as shown above
-def initialize_llmchain(vectordb, model_name="gpt-3.5-turbo", temperature=0.9):
-    chat_model = ChatOpenAI(temperature=temperature, model_name=model_name)
-    chat_vector_db_chain = ChatVectorDBChain.from_llm(
-        llm=chat_model,
-        vectordb=vectordb,
-        return_source_documents=True
+    else:
+        llm = HuggingFaceEndpoint(
+            repo_id=llm_model, 
+            # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "trust_remote_code": True, "torch_dtype": "auto"}
+            # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k}
+            temperature = temperature,
+            max_new_tokens = max_tokens,
+            top_k = top_k,
+        )
+    
+    progress(0.75, desc="Defining buffer memory...")
+    memory = ConversationBufferMemory(
+        memory_key="chat_history",
+        output_key='answer',
+        return_messages=True
     )
-    return chat_vector_db_chain
+    # retriever=vector_db.as_retriever(search_type="similarity", search_kwargs={'k': 3})
+    retriever=vector_db.as_retriever()
+    progress(0.8, desc="Defining retrieval chain...")
+    qa_chain = ConversationalRetrievalChain.from_llm(
+        llm,
+        retriever=retriever,
+        chain_type="stuff", 
+        memory=memory,
+        # combine_docs_chain_kwargs={"prompt": your_prompt})
+        return_source_documents=True,
+        #return_generated_question=False,
+        verbose=False,
+    )
+    progress(0.9, desc="Done!")
+    return qa_chain
 
-# Example usage
-if __name__ == "__main__":
-    # Suppose `documents` is a list of your documents and `collection_name` is your desired collection name
-    documents = [...]  # your documents here
-    collection_name = "your_collection_name"
-    vectordb = create_db(documents, collection_name)
 
-    # Initialize the ChatVectorDBChain with the vector database
-    chat_vector_db_chain = initialize_chat_vector_db_chain(vectordb)
-    
-    # Use the chain to process a query
-    query = "your query here"
-    result = chat_vector_db_chain({"question": query, "chat_history": ""})
-    
-    print("Answer:")
-    print(result["answer"])
 # Generate collection name for vector database
 #  - Use filepath as input, ensuring unicode text
 def create_collection_name(filepath):
@@ -155,31 +208,36 @@ def format_chat_history(message, chat_history):
     return formatted_chat_history
     
 
-def conversation(qa_chain, message, history):
+def conversation(vector_db, message, history):
+    # Initialize ChatOpenAI with your desired settings
+    chat_open_ai = ChatOpenAI(temperature=0.9, model_name="gpt-3.5-turbo")
+    # Assuming vector_db is already initialized correctly
+    pdf_ga = ChatVectorDBChain.from_llm(chat_open_ai, vector_db, return_source_documents=True)
+    
+    # Format the chat history for input to the model
     formatted_chat_history = format_chat_history(message, history)
-    #print("formatted_chat_history",formatted_chat_history)
-   
+    
     # Generate response using QA chain
-    response = qa_chain({"question": message, "chat_history": formatted_chat_history})
-    response_answer = response["answer"]
-    if response_answer.find("Helpful Answer:") != -1:
+    result = pdf_ga({"question": message, "chat_history": formatted_chat_history})
+    response_answer = result["answer"]
+    if "Helpful Answer:" in response_answer:
         response_answer = response_answer.split("Helpful Answer:")[-1]
-    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()
-    # Langchain sources are zero-based
-    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
-    # print ('chat response: ', response_answer)
-    # print('DB source', response_sources)
+    
+    response_sources = result["source_documents"]
+    response_source1 = response_sources[0].page_content.strip() if len(response_sources) > 0 else "No source available"
+    response_source2 = response_sources[1].page_content.strip() if len(response_sources) > 1 else "No source available"
+    response_source3 = response_sources[2].page_content.strip() if len(response_sources) > 2 else "No source available"
+    
+    # Langchain sources are zero-based; adjust pages if sources are available
+    response_source1_page = response_sources[0].metadata["page"] + 1 if len(response_sources) > 0 else "N/A"
+    response_source2_page = response_sources[1].metadata["page"] + 1 if len(response_sources) > 1 else "N/A"
+    response_source3_page = response_sources[2].metadata["page"] + 1 if len(response_sources) > 2 else "N/A"
     
     # Append user message and response to chat history
     new_history = history + [(message, response_answer)]
-    # return gr.update(value=""), new_history, response_sources[0], response_sources[1] 
-    return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
     
+    return gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
+
 
 def upload_file(file_obj):
     list_file_path = []
@@ -190,7 +248,7 @@ def upload_file(file_obj):
     # initialize_database(file_path, progress)
     return list_file_path
 
-        
+
 def demo():
     with gr.Blocks(theme="base") as demo:
         vector_db = gr.State()
@@ -201,7 +259,6 @@ def demo():
         """PDF-based chatbot (by Dr. Aloke Upadhaya)</center></h2>
         <h3>Ask any questions about your PDF documents, along with follow-ups</h3>
         """)
-            
         with gr.Tab("Step 1 - Document pre-processing"):
             with gr.Row():
                 document = gr.Files(height=100, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload your PDF documents (single or multiple)")
@@ -282,3 +339,4 @@ def demo():
 
 if __name__ == "__main__":
     demo()
+