Adding chat

app.py

@@ -6,6 +6,7 @@ from llama_index.core import Settings
 from llama_index.llms.gemini import Gemini
 from llama_index.core import Document, VectorStoreIndex
 from llama_index.embeddings.gemini import GeminiEmbedding
+from llama_index.core import load_index_from_storage, StorageContext
 
 reader = easyocr.Reader(['en'])
 
@@ -16,23 +17,32 @@ gemini_embedding_model = GeminiEmbedding(api_key=os.getenv('GEMINI_API_KEY'), mo
 Settings.llm = llm
 Settings.embed_model = gemini_embedding_model
 
-def inference(img_path, width_ths):
+def ocr_inference(img_path, width_ths):
     output = reader.readtext(img_path, detail=0, slope_ths=0.7, ycenter_ths=0.9,
                           height_ths=0.8, width_ths=width_ths, add_margin=0.2)
     
     output = "\n".join(output)
     
-    # create a Document object from the extracted text
     doc = Document(text = output)
 
-    # Create an index from the documents and save it to the disk.
     index = VectorStoreIndex.from_documents([doc])
 
-    # save the index
     index.storage_context.persist(persist_dir = "./receiptsembeddings")
     
     return output
 
+def inference(question):
+    persist_dir = "./receiptsembeddings"
+
+    storage_context = StorageContext.from_defaults(persist_dir = persist_dir)
+    index = load_index_from_storage(storage_context)
+
+    query_engine = index.as_query_engine()
+
+    response = query_engine.query(question)
+
+    return response
+
 title = "Receipt RAG"
 description = "A simple Gradio interface to query receipts using RAG"
 examples = [["data/receipt_00000.JPG", 7.7],
@@ -50,17 +60,19 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column():
             ocr_out = gr.Textbox(label="OCR Output", type="text")
 
-    submit_btn.click(inference, inputs=[image, width_ths], outputs=ocr_out)
-    clear_btn.click(lambda: [None, 7.7], inputs=[image, width_ths])
+    submit_btn.click(ocr_inference, inputs=[image, width_ths], outputs=ocr_out)
 
     examples_obj = gr.Examples(examples=examples, inputs=[image, width_ths])
 
-# demo = gr.Interface(inference, 
-#                     inputs = [gr.Image(width=320, height=320, label="Input Receipt"), 
-#                               gr.Slider(0, 10, 7.7, 0.1, label="Width Threshold to merge bounding boxes")],
-#                     outputs= [gr.Textbox(label="OCR Output", type="text")],
-#                     title=title,
-#                     description=description,
-#                     examples=examples)
+    with gr.Row():
+        with gr.Column():
+            text = gr.Textbox(label="Question", type="text")
+            with gr.Row():
+                chat_clear_btn = gr.ClearButton(components=[text])
+                chat_submit_btn = gr.Button("Submit", variant='primary')
+        with gr.Column():
+            chat_out = gr.Textbox(label="Response", type="text")
+    
+    chat_clear_btn.click(inference, inputs=[text], outputs=[chat_out])
 
 demo.launch()