Update app.py

app.py

@@ -1,38 +1,36 @@
 import os
 import time
-from transformers import pipeline
-from langchain.embeddings import HuggingFaceEmbeddings
-from huggingface_hub import InferenceClient
+
 import langchain
-from langchain import HuggingFaceHub
-from langchain.cache import InMemoryCache
+# loaders
 from langchain.document_loaders import PyPDFLoader, OnlinePDFLoader, Docx2txtLoader, UnstructuredWordDocumentLoader, UnstructuredPowerPointLoader
-from langchain.text_splitter import CharacterTextSplitter,RecursiveCharacterTextSplitter
+# splits
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+# embeddings
+from langchain.embeddings import HuggingFaceEmbeddings, OpenAIEmbeddings
+# vector stores
 from langchain.vectorstores import Chroma
+# huggingface hub
+from huggingface_hub import InferenceClient
+from langchain import HuggingFaceHub
+# models
+from langchain.llms import OpenAI
+# retrievers
 from langchain.chains import RetrievalQA
 import gradio as gr
 
-def define_embeddings_llm(openai_key):
+HUGGINGFACEHUB_API_TOKEN = os.environ["HUGGINGFACEHUB_API_TOKEN"]
+
+def build_context(openai_key,files,urls):
   if openai_key != "":
-    embeddings = OpenAIEmbeddings(openai_api_key=openai_key)
-    llm = OpenAI(
-        temperature=0, model_name="gpt-3.5-turbo-16k", openai_api_key=openai_key, verbose=False
-    )
+    embeddings = OpenAIEmbeddings(model_name="text-embedding-ada-002", openai_api_key=openai_key)
   else:
-    HUGGINGFACEHUB_API_TOKEN = os.environ["HUGGINGFACEHUB_API_TOKEN"]
     embeddings = HuggingFaceEmbeddings(
-                model_name="sentence-transformers/all-MiniLM-L6-v2"
+                model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={'device': 'cpu'}
             )
-    llm = HuggingFaceHub(repo_id="MBZUAI/LaMini-Flan-T5-248M" ,
-                         model_kwargs={"max_length":2048,"do_sample":True,
-                                       "temperature":0.2})
-  langchain.llm_cache = InMemoryCache()
-  return embeddings,llm
-
-def build_context(openai_key,files,urls):
-  embeddings, llm = define_embeddings_llm(openai_key)
   documents = []
   if files is not None:
+    print("files not none")
     for idx, file in enumerate(files):
       if file.name.endswith('.pdf'):
         loader = PyPDFLoader(file.name)
@@ -44,63 +42,71 @@ def build_context(openai_key,files,urls):
         loader = UnstructuredPowerPointLoader(file.name)
         documents.extend(loader.load())
   if urls != "":
+    print("urls not none")
     list_urls = urls.split(sep=",")
-    for url in list_urls: 
+    for url in list_urls:
       loader = OnlinePDFLoader(url)
       documents.extend(loader.load())
-  #text_splitter = RecursiveCharacterTextSplitter(chunk_size=400,chunk_overlap=20,length_function=len,separators=["\n\n", "\n", " ", ""])
-  text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+  text_splitter = RecursiveCharacterTextSplitter(chunk_size=800,chunk_overlap=0,length_function=len,separators=["\n\n", "\n", " ", ""])
   chunked_documents = text_splitter.split_documents(documents)
   global vectordb
   vectordb = Chroma.from_documents(
     documents=chunked_documents,
-    embedding=embeddings
-  )
-  global qa_chain
-  qa_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    retriever=vectordb.as_retriever(search_kwargs={'k': 7}),
-    chain_type="stuff",
-    #return_source_documents=True
-  )
-  return "ready"
-
-def loading():
-    return "Loading..."
+    embedding=embeddings,
+    )
+  return "loaded"
 
-def respond(message, chat_history):
-  result = qa_chain({"query": message})["result"]
+def llm_response(openai_key, message, chat_history):
+  if openai_key != "":
+    llm = OpenAI(
+        temperature=0, openai_api_key=openai_key, model_name="gpt-3.5-turbo", verbose=False
+    )
+  else:
+    llm = HuggingFaceHub(repo_id='MBZUAI/LaMini-Flan-T5-248M',
+                        huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
+                         model_kwargs={"max_length":512,"do_sample":True,
+                                       "temperature":0.2})
+  qa_chain = RetrievalQA.from_chain_type(llm = llm,
+                                       chain_type = "stuff", # map_reduce, map_rerank, stuff, refine
+                                       retriever = vectordb.as_retriever(search_kwargs = {"k": 10}), 
+                                       #chain_type_kwargs = {"prompt": PROMPT},
+                                       return_source_documents = False,
+                                       verbose = True)
+  result = qa_chain(message)["result"]
   chat_history.append((message, result))
   time.sleep(2)
   return "", chat_history
+  
+def loading():
+    return "Loading..."
 
 def clear_chromadb():
   ids = vectordb.get()["ids"]
   for id in ids:
     vectordb._collection.delete(ids=id)
 
-with gr.Blocks() as demo:
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
   with gr.Row():
     openai_key = gr.Textbox(label="Enter your OpenAI API Key if you want to use the gpt-3.5-turbo-16k model. If not, the open source LaMini-Flan-T5-248M is used")
   with gr.Row():
-    with gr.Column():
-      pdf_docs = gr.Files(label="Load pdf/docx/ppt/pptx files", file_types=['.pdf','.docx','.ppt','.pptx'], type="file")
-    with gr.Column():
-      urls = gr.Textbox(label="Enter one of multiple online pdf urls (comma separated if multiple)", value=None)
+      pdf_docs = gr.Files(label="Load pdf files", file_types=['.pdf','.docx','.ppt','.pptx'], type="file")
+      urls = gr.Textbox(label="Enter one of multiple online pdf urls (comma separated if multiple)")
   with gr.Row():
-    load_context = gr.Button("Load documents and urls")
+    load_docs = gr.Button("Load documents and urls", variant="primary", scale=1)
+    loading_status = gr.Textbox(label="Loading status", placeholder="", interactive=False, scale=0)
+    if loading_status == "loaded":
+      gr.Info("Documents loaded")
   with gr.Row():
-    loading_status = gr.Textbox(label="Status", placeholder="", interactive=False)
-  with gr.Row():
-    with gr.Column():
-      hg_chatbot = gr.Chatbot()
+    with gr.Column(scale=1):
       msg = gr.Textbox(label="User message")
-      clear = gr.ClearButton([msg, hg_chatbot])
-      cleardb = gr.Button(value="Réinitialiser le contexte")
-  load_context.click(loading, None, loading_status, queue=False)
-  load_context.click(build_context, inputs=[openai_key,pdf_docs, urls], outputs=[loading_status], queue=False)
-  msg.submit(respond, [msg, hg_chatbot], [msg, hg_chatbot])
+      chatbot = gr.Chatbot()
+  with gr.Row():
+      clearchat = gr.ClearButton([msg, chatbot], value="New chat",)
+      cleardb = gr.Button(value="Reset context (for loading new documents)", variant="secondary")
+  load_docs.click(loading, None, loading_status, queue=False)
+  load_docs.click(build_context, inputs=[openai_key,pdf_docs, urls], outputs=[loading_status], queue=False)
+  msg.submit(llm_response, [openai_key, msg, chatbot], [msg, chatbot])
   cleardb.click(clear_chromadb)
-
+    
 demo.queue(concurrency_count=3)
 demo.launch()