Update app.py

app.py

@@ -14,22 +14,27 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import transformers
 import torch
 
-
+# load the model
 model_name = 'google/flan-t5-base'
-model = T5ForConditionalGeneration.from_pretrained(model_name, device_map='auto', offload_folder="offload")
+model = T5ForConditionalGeneration.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
-
+# to calculate text embeddings
 ST_name = 'sentence-transformers/sentence-t5-base'
 st_model = SentenceTransformer(ST_name)
-print('sentence read')
 
+# to store our embeddings and search
 client = chromadb.Client()
-collection = client.create_collection("test_db") 
+collection = client.create_collection("my_db") 
 
 
 def get_context(query_text):
+    '''
+    Given query in tokenized format, find its embeddings
+    Search in Chroma DB
+    and return results
+    '''
     
     query_emb = st_model.encode(query_text)
     query_response = collection.query(query_embeddings=query_emb.tolist(), n_results=4)
@@ -37,13 +42,23 @@ def get_context(query_text):
     context = context.replace('\n', ' ').replace('  ', ' ')
     return context
 
+
+
 def local_query(query, context):
-    t5query = """Using the available context, please answer the question. 
-    If you aren't sure please say i don't know.
+    '''
+    Given query (user response)
+    Construct LLM query adding context to it
+    Return response of LLM
+    '''
+
+    
+    t5query = """Please answer the question based on the given context. 
+    If you are not sure about your response, say I am not sure.
     Context: {}
     Question: {}
     """.format(context, query)
 
+    # calculate embeddings for the query
     inputs = tokenizer(t5query, return_tensors="pt")
 
     outputs = model.generate(**inputs, max_new_tokens=20)
@@ -55,48 +70,59 @@ def local_query(query, context):
    
    
 
-def run_query(btn, history, query):
+def run_query(history, query):
+    '''
+    Run Gradio ChatInterface
+    Given user response (query), find the most similar/related part to the question from the uploaded document
+    Using Chroma search
+    Update the query with context, and ask the question to LLM
+    '''
 
-    context = get_context(query)
-    
-    print('calling local query')
-    result = local_query(query, context)
-
-    
-    print('printing result after call back')
-    print(result)
+    context = get_context(query) # find the related part from the pdf
+    result = local_query(query, context) #  add context to model query
+ 
 
-    history.append((query, str(result[0])))
+    history.append((query, str(result[0])))  # append result to chatInterface history
         
 
-    print('printing history')
-    print(history)
     return  history, ""
 
 
 
 def upload_pdf(file):
+    '''
+    Upload a PDF
+    Split into chunks 
+    Encode each chunk into embeddings
+    Assign a unique ID for each chunk embedding
+    Construct Chroma DB
+    Update your global Chroma DB collection
+    '''
     try:
         if file is not None: 
 
             global collection
             
             file_name = file.name 
-   
+
+            #  Upload pdf document
             loader = PDFMinerLoader(file_name)
             doc = loader.load()
-        
-            text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+
+            #  extract chunks
+            text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=10)
             texts = text_splitter.split_documents(doc)
         
             texts = [i.page_content for i in texts]
-        
+
+            #  find embedding for each chunk
             doc_emb = st_model.encode(texts)
             doc_emb = doc_emb.tolist()
-        
+
+            #  index the embeddings
             ids = [str(uuid.uuid1()) for _ in doc_emb]
         
-            
+            #  add each chunk embedding to ChromaDB
             collection.add(
                 embeddings=doc_emb,
                 documents=texts,
@@ -116,26 +142,28 @@ def upload_pdf(file):
     
  
 with gr.Blocks() as demo:  
-    
+    '''
+    Frontend for our tool
+    '''
+
+    #  Upload a PDF focument
     btn = gr.UploadButton("Upload a PDF", file_types=[".pdf"])
-    output = gr.Textbox(label="Output Box")
-    chatbot = gr.Chatbot(height=240)
+    output = gr.Textbox(label="Output Box") #  to put message indicating the status of upload
+    chatbot = gr.Chatbot(height=240) #  our chatbot interface
     
     with gr.Row():
         with gr.Column(scale=0.70):
             txt = gr.Textbox(
                 show_label=False,
-                placeholder="Enter a question",
+                placeholder="Type a question",
             ) 
 
- 
-    # Event handler for uploading a PDF
+
+    # Backend for our tool
+    # Event handlers
     btn.upload(fn=upload_pdf, inputs=[btn], outputs=[output])
-    txt.submit(run_query, [btn, chatbot, txt], [chatbot, txt])
-    #.then(
-            # generate_response, inputs =[chatbot,],outputs = chatbot,)
+    txt.submit(run_query, [chatbot, txt], [chatbot, txt])
 
 
 gr.close_all()
-# demo.launch(share=True)
-demo.queue().launch()
+demo.queue().launch() # use query for a better performance