Update app.py

app.py

@@ -6,10 +6,7 @@ from PyPDF2 import PdfReader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.callbacks.manager import CallbackManager
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-from langchain.llms import LlamaCpp 
-
-from langchain.vectorstores import Qdrant
-from transformers import AutoModelForCausalLM
+from ctransformers import AutoModelForCausalLM
 
 # Load the embedding model
 encoder = SentenceTransformer('jinaai/jina-embedding-b-en-v1')
@@ -17,8 +14,6 @@ print("Embedding model loaded...")
 
 # Load the LLM
 callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
-
-'''
 llm = AutoModelForCausalLM.from_pretrained(
     "TheBloke/Llama-2-7B-Chat-GGUF",
     model_file="llama-2-7b-chat.Q3_K_S.gguf",
@@ -27,32 +22,25 @@ llm = AutoModelForCausalLM.from_pretrained(
     repetition_penalty=1.5,
     max_new_tokens=300,
 )
-'''
-llm = LlamaCpp(
-    model_path="./llama-2-7b-chat.Q3_K_S.gguf",
-    temperature = 0.2,
-    n_ctx=2048,
-    f16_kv=True,  # MUST set to True, otherwise you will run into problem after a couple of calls
-    max_tokens = 500,
-    callback_manager=callback_manager,
-    verbose=True,
-)
 print("LLM loaded...")
 
-client = QdrantClient(path="./db")
+def get_chunks(text):
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=250,
+        chunk_overlap=50,
+        length_function=len,
+    )
+    return text_splitter.split_text(text)
 
 def setup_database(files):
     all_chunks = []
     for file in files:
-        pdf_path = file
-        reader = PdfReader(pdf_path)
-        text = "".join(page.extract_text() for page in reader.pages if page.extract_text())
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=250, chunk_overlap=50, length_function=len)
-        chunks = text_splitter.split_text(text)
+        reader = PdfReader(file)
+        text = "".join(page.extract_text() for page in reader.pages)
+        chunks = get_chunks(text)
         all_chunks.extend(chunks)
-    
-    print(f"Total chunks: {len(all_chunks)}")
-    
+
+    client = QdrantClient(path="./db")
     client.recreate_collection(
         collection_name="my_facts",
         vectors_config=models.VectorParams(
@@ -60,51 +48,64 @@ def setup_database(files):
             distance=models.Distance.COSINE,
         ),
     )
-    
-    print("Collection created...")
 
-    for idx, chunk in enumerate(all_chunks):
-        client.upload_record(
-            collection_name="my_facts",
-            record=models.Record(
-                id=idx,
-                vector=encoder.encode(chunk).tolist(),
-                payload={"text": chunk}
-            )
-        )
+    records = [
+        models.Record(
+            id=idx,
+            vector=encoder.encode(chunk).tolist(),
+            payload={f"chunk_{idx}": chunk}
+        ) for idx, chunk in enumerate(all_chunks)
+    ]
     
-    print("Records uploaded...")
+    client.upload_records(
+        collection_name="my_facts",
+        records=records,
+    )
 
-def answer(question):
+def answer_question(question):
+    client = QdrantClient(path="./db")
     hits = client.search(
         collection_name="my_facts",
         query_vector=encoder.encode(question).tolist(),
         limit=3
     )
     
-    context = " ".join(hit.payload["text"] for hit in hits)
-    system_prompt = "You are a helpful co-worker. Use the provided context to answer user questions. Do not use any other information."
-    prompt = f"Context: {context}\nUser: {question}\n{system_prompt}"
-    response = llm(prompt)
+    context = " ".join(hit.payload[f"chunk_{hit.id}"] for hit in hits)
+
+    system_prompt = """You are a helpful co-worker, you will use the provided context to answer user questions.
+    Read the given context before answering questions and think step by step. If you cannot answer a user question based on 
+    the provided context, inform the user. Do not use any other information for answering user. Provide a detailed answer to the question."""
+
+    B_INST, E_INST = "[INST]", "[/INST]"
+    B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
+
+    instruction = f"Context: {context}\nUser: {question}"
+    prompt_template = f"{B_INST}{B_SYS}{system_prompt}{E_SYS}{instruction}{E_INST}"
+    
+    response = llm(prompt_template)
     return response
 
-def chat(messages):
-    if not messages:
-        return "Please upload PDF documents to initialize the database."
-    last_message = messages[-1]
-    return answer(last_message["message"])
+def chat(messages, files):
+    if files:
+        setup_database(files)
+    if messages:
+        question = messages[-1]["text"]
+        answer = answer_question(question)
+        messages.append({"text": answer, "is_user": False})
+    return messages
 
-screen = gr.Interface(
+interface = gr.Interface(
     fn=chat,
-    inputs=gr.Textbox(placeholder="Type your question here..."),
-    outputs="chatbot",
+    inputs=[
+        gr.Chatbot(label="Chat"),
+        gr.File(label="Upload PDFs", file_count="multiple")
+    ],
+    outputs=gr.Chatbot(label="Chat"),
     title="Q&A with PDFs 👩🏻‍💻📓✍🏻💡",
     description="This app facilitates a conversation with PDFs uploaded💡",
     theme="soft",
+    share=True,
     live=True,
-    allow_flagging=False,
 )
 
-
-# Add a way to upload and setup the database before starting the chat
-screen.launch()
+interface.launch()