Updated code files

app.py

@@ -3,13 +3,13 @@ from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoPro
 import torch
 from qwen_vl_utils import process_vision_info
 from PIL import Image
-import re
 import gradio as gr
+import re
 
 rag = RAGMultiModalModel.from_pretrained("vidore/colpali")
 vlm = Qwen2VLForConditionalGeneration.from_pretrained(
     "Qwen/Qwen2-VL-2B-Instruct",
-    torch_dtype=torch.float32,
+    torch_dtype=torch.float16,
     trust_remote_code=True,
     device_map="auto",
 )
@@ -26,59 +26,58 @@ def extract_text(image, query):
             ],
         }
     ]
-    
+
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    
+
     image_inputs, video_inputs = process_vision_info(messages)
     inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt")
-    inputs = inputs.to("cpu")
-    with torch.no_grad():
-        generated_ids = vlm.generate(**inputs, max_new_tokens=200, temperature=0.7, top_p=0.9)
-        generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
-        return processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    
-
-def ocr(image):
-    queries = [
-        # "Extract and transcribe all the text visible in the image, including any small or partially visible text.",
-        # "Look closely at the image and list any text you see, no matter how small or unclear.",
-        # "What text can you identify in this image? Include everything, even if it's partially obscured or in the background."
-        "Extract all the text in Sanskrit and English from the image."
-    ]
+    inputs = inputs.to("cuda")
 
-    all_extracted_text = []
-    for query in queries:
-        extracted_text = extract_text(image, query)
-        all_extracted_text.append(extracted_text)
+    generated_ids = vlm.generate(**inputs, max_new_tokens=200, temperature=0.7, top_p=0.9)
+    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    return processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
 
-    # Combine and deduplicate the results
-    final_text = "\n".join(set(all_extracted_text))
-
-    # final_text = post_process_text(final_text)
-    return final_text
-    
-    
-def main_fun(image, keyword):
-    ext_text = ocr(image)
-    
-    if keyword:
-        highlight_text = re.sub(f'({re.escape(keyword)})', r'<span style="background-color: yellow;">\1</span>', ext_text, flags=re.IGNORECASE)
+def search_text(text, query):
+    if query:
+        searched_text = re.sub(f'({re.escape(query)})', r'<span style="background-color: yellow;">\1</span>', text, flags=re.IGNORECASE)
     else:
-        highlight_text = ext_text
-    
-    return ext_text, highlight_text
+        searched_text = text
+    return searched_text
 
-iface = gr.Interface(
-    fn=main_fun,
-    inputs=[
-        gr.Image(type="pil", label="Upload an Image"),
-        gr.Textbox(label="Enter search term", placeholder="Search")
-    ],
-    outputs=[
-        gr.Textbox(label="Extracted Text"),
-        gr.HTML(label="Search Results")
-    ],
-    title="Document Search using OCR (English/Hindi)"
-)
+def extraction(image, query):
+    extracted_text = extract_text(image, query)
+    return extracted_text, extracted_text   # return twice - one to display output and the other for state management
 
-iface.launch()
+
+"""
+    Main App
+"""
+with gr.Blocks() as main_app:
+    gr.Markdown("# Document Reader using OCR(English/Hindi)")
+    
+    with gr.Row():
+        with gr.Column():
+            img_input = gr.Image(type="pil", label="Upload an Image")
+            query_input = gr.Textbox(label="Enter query for retrieval", placeholder="Query/Prompt")
+            search_input = gr.Textbox(label="Enter search term", placeholder="Search")
+            extract_button = gr.Button("Read Doc!")
+            search_button = gr.Button("Search!")
+            
+        with gr.Column():
+            extracted_text_op = gr.Textbox(label="Output")
+            search_text_op = gr.HTML(label="Search Results")
+        
+        extracted_text_state = gr.State()
+        extract_button.click(
+            extraction,
+            inputs=[img_input, query_input],
+            outputs=[extracted_text_op, extracted_text_state]
+        )
+        
+        search_button.click(
+            search_text,
+            inputs=[extracted_text_state, search_input],
+            outputs=[search_text_op]
+        )
+        
+main_app.launch()