Update app.py

app.py

@@ -1,110 +1,110 @@
-import gradio as gr
-from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
-from PIL import Image
-import torch
-import re
-
-# Load the pre-trained model and processor
-model = Qwen2VLForConditionalGeneration.from_pretrained(
-    "Qwen/Qwen2-VL-2B-Instruct",
-    torch_dtype="auto",
-    device_map="auto",
-)
-
-processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
-
-# Function to extract text from the image
-def extract_text(image):
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {"type": "text", "text": "can u extract the text in hindi"}
-            ]
-        }
-    ]
-    
-    # Process input image and text prompt
-    text_prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-    
-    inputs = processor(
-        text=[text_prompt],
-        images=[image],
-        padding=True,
-        return_tensors="pt"
-    )
-    
-    inputs = inputs.to("cuda" if torch.cuda.is_available() else "cpu")
-    
-    # Generate output text from the model
-    output_ids = model.generate(**inputs, max_new_tokens=1024)
-    
-    generated_ids = [
-        output_ids[len(input_ids):]
-        for input_ids, output_ids in zip(inputs.input_ids, output_ids)
-    ]
-    
-    # Decode the generated text
-    extracted_text = processor.batch_decode(
-        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
-    )[0]  # Extracted text
-    
-    return extracted_text
-
-# Function to highlight keywords in the text, even for right-to-left scripts like Hindi
-def highlight_keywords(extracted_text, keywords):
-    highlighted_text = extracted_text
-    if keywords:
-        for keyword in keywords.split(","):
-            keyword = keyword.strip()
-            if keyword:
-                # Ensure correct Unicode support for keywords (use re.UNICODE for non-ASCII)
-                highlighted_text = re.sub(
-                    re.escape(keyword),  # Use re.escape to handle special characters in keywords
-                    r'<mark>\g<0></mark>',  # Highlight the found keyword
-                    highlighted_text,
-                    flags=re.IGNORECASE | re.UNICODE  # Ignore case, and handle Unicode characters
-                )
-    
-    return highlighted_text
-
-# First step: Extract text from the uploaded image
-def extract_text_step(image):
-    extracted_text = extract_text(image)
-    return extracted_text, extracted_text  # Return extracted text and store it in state
-
-# Second step: Search and highlight keywords in the extracted text
-def highlight_keywords_step(extracted_text, keywords):
-    highlighted_text = highlight_keywords(extracted_text, keywords)
-    return highlighted_text
-
-# Gradio UI
-with gr.Blocks() as demo:
-    # Step 1: Image Upload and Text Extraction
-    with gr.Row():
-        image_input = gr.Image(type="pil", label="Upload Image")
-        extract_button = gr.Button("Extract Text")
-    extracted_text_output = gr.Textbox(label="Extracted Text")
-    
-    # Step 2: Keyword Input and Highlighting
-    with gr.Row():
-        keyword_input = gr.Textbox(label="Enter keywords (comma-separated)", placeholder="Enter keywords after text extraction")
-        search_button = gr.Button("Highlight Keywords")
-    highlighted_text_output = gr.HTML(label="Highlighted Text with Matches")
-    
-    # Define interactions
-    extract_button.click(
-        fn=extract_text_step,  # Call text extraction function
-        inputs=image_input,
-        outputs=[extracted_text_output, extracted_text_output],  # Display text and store in state
-    )
-    
-    search_button.click(
-        fn=highlight_keywords_step,  # Call keyword highlighting function
-        inputs=[extracted_text_output, keyword_input],  # Use extracted text and keywords
-        outputs=highlighted_text_output,  # Display highlighted text
-    )
-
-# Launch the app
-demo.launch()
+import gradio as gr
+from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
+from PIL import Image
+import torch
+import re
+
+# Load the pre-trained model and processor
+model = Qwen2VLForConditionalGeneration.from_pretrained(
+    "Qwen/Qwen2-VL-2B-Instruct",
+    torch_dtype="auto",
+    device_map="auto",
+)
+
+processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
+
+# Function to extract text from the image
+def extract_text(image):
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": "can u extract the text in hindi"}
+            ]
+        }
+    ]
+    
+    # Process input image and text prompt
+    text_prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    
+    inputs = processor(
+        text=[text_prompt],
+        images=[image],
+        padding=True,
+        return_tensors="pt"
+    )
+    
+    inputs = inputs.to("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Generate output text from the model
+    output_ids = model.generate(**inputs, max_new_tokens=1024)
+    
+    generated_ids = [
+        output_ids[len(input_ids):]
+        for input_ids, output_ids in zip(inputs.input_ids, output_ids)
+    ]
+    
+    # Decode the generated text
+    extracted_text = processor.batch_decode(
+        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+    )[0]  # Extracted text
+    
+    return extracted_text
+
+# Function to highlight keywords in the text, even for right-to-left scripts like Hindi
+def highlight_keywords(extracted_text, keywords):
+    highlighted_text = extracted_text
+    if keywords:
+        for keyword in keywords.split(","):
+            keyword = keyword.strip()
+            if keyword:
+                # Ensure correct Unicode support for keywords (use re.UNICODE for non-ASCII)
+                highlighted_text = re.sub(
+                    re.escape(keyword),  # Use re.escape to handle special characters in keywords
+                    r'<mark>\g<0></mark>',  # Highlight the found keyword
+                    highlighted_text,
+                    flags=re.IGNORECASE | re.UNICODE  # Ignore case, and handle Unicode characters
+                )
+    
+    return highlighted_text
+
+# First step: Extract text from the uploaded image
+def extract_text_step(image):
+    extracted_text = extract_text(image)
+    return extracted_text, extracted_text  # Return extracted text and store it in state
+
+# Second step: Search and highlight keywords in the extracted text
+def highlight_keywords_step(extracted_text, keywords):
+    highlighted_text = highlight_keywords(extracted_text, keywords)
+    return highlighted_text
+
+# Gradio UI
+with gr.Blocks() as demo:
+    # Step 1: Image Upload and Text Extraction
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload Image")
+        extract_button = gr.Button("Extract Text")
+    extracted_text_output = gr.Textbox(label="Extracted Text")
+    
+    # Step 2: Keyword Input and Highlighting
+    with gr.Row():
+        keyword_input = gr.Textbox(label="Enter keywords (comma-separated)", placeholder="Enter keywords after text extraction")
+        search_button = gr.Button("Highlight Keywords")
+    highlighted_text_output = gr.HTML(label="Highlighted Text with Matches")
+    
+    # Define interactions
+    extract_button.click(
+        fn=extract_text_step,  # Call text extraction function
+        inputs=image_input,
+        outputs=[extracted_text_output, extracted_text_output],  # Display text and store in state
+    )
+    
+    search_button.click(
+        fn=highlight_keywords_step,  # Call keyword highlighting function
+        inputs=[extracted_text_output, keyword_input],  # Use extracted text and keywords
+        outputs=highlighted_text_output,  # Display highlighted text
+    )
+
+# Launch the app
+demo.launch()