Update chat.py

chat.py

@@ -1,63 +1,40 @@
 import os
 import openai
-import torch
-
-import llama_index
-from llama_index import ServiceContext, set_global_service_context
-from llama_index.prompts import PromptTemplate
-from llama_index.llms import HuggingFaceLLM
-from llama_index import VectorStoreIndex, SimpleDirectoryReader
+import langchain
+from langchain.document_loaders import PyMuPDFLoader, DirectoryLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.chat_models import ChatOpenAI
+from langchain.chains import RetrievalQA
+from langchain.globals import set_verbose
+from langchain.llms import VLLM
 
 
+class RetrievalChatbot():
+    def __init__(self, api_key, api_base, model_name):
+        openai.api_key = api_key
+        #openai.api_base = api_base
+        self.model_name = model_name
 
-def messages_to_prompt(messages):
-  prompt = ""
-  for message in messages:
-    if message.role == 'system':
-      prompt += f"<|system|>\n{message.content}</s>\n"
-    elif message.role == 'user':
-      prompt += f"<|user|>\n{message.content}</s>\n"
-    elif message.role == 'assistant':
-      prompt += f"<|assistant|>\n{message.content}</s>\n"
+        set_verbose(True)
 
-  # ensure we start with a system prompt, insert blank if needed
-  if not prompt.startswith("<|system|>\n"):
-    prompt = "<|system|>\n</s>\n" + prompt
 
-  # add final assistant prompt
-  prompt = prompt + "<|assistant|>\n"
 
-  return prompt
+        embeddings = OpenAIEmbeddings(max_retries=100)
 
+        if os.path.exists("persist"):
+            vectordb = Chroma(persist_directory="persist", embedding_function=embeddings)
+            print("loaded existing database")
+        else:
+            os.mkdir("persist")
 
-class RetrievalChatbot():
-    def __init__(self, api_key, api_base, pdf_dir, model_name):
-        openai.api_key = api_key
-        openai.api_base = api_base
-        self.model_name = model_name
-        documents = SimpleDirectoryReader(input_dir="papers_all").load_data()
-        print("find doc")
-        llm_zephyr = HuggingFaceLLM(
-            model_name="HuggingFaceH4/zephyr-7b-beta",
-            tokenizer_name="HuggingFaceH4/zephyr-7b-beta",
-            query_wrapper_prompt=PromptTemplate("<|system|>\n</s>\n<|user|>\n{query_str}</s>\n<|assistant|>\n"),
-            context_window=2048,
-            max_new_tokens=128,
-            messages_to_prompt=messages_to_prompt,
-            device_map="auto",
+        retriever = vectordb.as_retriever(
+            search_type="mmr",
+            search_kwargs={"k": 10, "fetch_k": 50}
         )
-
-        print("loaded llm")
-        service_context = ServiceContext.from_defaults(llm=llm_zephyr, chunk_size=512)
-    # set_global_service_context(service_context)
-        print("loaded doc")
-        index = VectorStoreIndex.from_documents(documents, service_context=service_context)
-       # index.storage_context.persist(persist_dir="index")
-        print("save index")
-        qa = index.as_query_engine(streaming=True)
-        self.qa=qa
-
-
+        llm = ChatOpenAI(model_name="gpt-4-1106-preview")
+        self.qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever)
 
         self.prompt = (
             "Please answer the following question using information with the assistance of the given context.\n",
@@ -98,10 +75,8 @@ class RetrievalChatbot():
         
         for rephrase in subquestion_list:
             query = "".join(self.prompt).format(message=rephrase)
-            print("query:",query)
-            response = self.qa.query(query)
-            print("respponse:",response)
-            responses_list.append(response)
+            retrieval_response = self.qa(query)["result"]
+            responses_list.append(retrieval_response)
 
         summarize_prompt = "".join(self.summarize_prompt).format(question=message, answers=responses_list)
         summarized_answer = self.get_openai_response(summarize_prompt, self.model_name)