Update app.py

app.py

@@ -5,6 +5,7 @@ from byaldi import RAGMultiModalModel
 from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
 from qwen_vl_utils import process_vision_info
 import torch
+import re
 
 # Load models
 def load_models():
@@ -59,8 +60,26 @@ def ocr_and_search(image, keyword):
     keyword_lower = keyword.lower()
     sentences = extracted_text.split('. ')
     matched_sentences = [sentence for sentence in sentences if keyword_lower in sentence.lower()]
+
+
+    #########################################################################################################################################################
     
-    return extracted_text, matched_sentences #, json_output
+    highlighted_text = matched_sentences
+    if search_queries:
+        keywords = [kw.strip() for kw in search_queries.split(",")]
+        for keyword in keywords:
+            highlighted_text = re.sub(
+                re.escape(keyword),
+                f"<span style='background-color: #f1c40f; color: #000; font-weight: bold;'>{keyword}</span>",
+                highlighted_text,
+                flags=re.IGNORECASE
+            )
+
+
+    ####################################################################################################################################################
+
+
+    return extracted_text, highlighted_text #, json_output
 
 
 # Gradio App