Update app.py

app.py

@@ -1,14 +1,20 @@
 import gradio as gr
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoProcessor, VisionEncoderDecoderModel
 from vllm import LLM, SamplingParams
+from PIL import Image
 
-# Load the model and tokenizer from Hugging Face
+# Load the language model and tokenizer from Hugging Face
 model_name = "facebook/opt-125m"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Initialize vLLM with CPU configuration
 vllm_model = LLM(model=model_name, tensor_parallel_size=1, device="cpu")
 
+# Load the OCR model and processor
+ocr_model_name = "microsoft/trocr-small-handwritten"
+ocr_model = VisionEncoderDecoderModel.from_pretrained(ocr_model_name)
+ocr_processor = AutoProcessor.from_pretrained(ocr_model_name)
+
 def generate_response(prompt, max_tokens, temperature, top_p):
     # Define sampling parameters
     sampling_params = SamplingParams(
@@ -24,53 +30,87 @@ def generate_response(prompt, max_tokens, temperature, top_p):
     generated_text = output[0].outputs[0].text
     return generated_text
 
+def ocr_image(image):
+    # Open the image and preprocess for OCR
+    image = Image.open(image)
+    pixel_values = ocr_processor(images=image, return_tensors="pt").pixel_values
+    
+    # Perform OCR
+    outputs = ocr_model.generate(pixel_values)
+    text = ocr_processor.decode(outputs[0], skip_special_tokens=True)
+    return text
+
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🚀 Hugging Face Integration with vLLM (CPU)")
-    gr.Markdown("Generate text using the vLLM integration with Hugging Face models on CPU.")
+    gr.Markdown("# 🚀 Hugging Face Integration with vLLM and OCR (CPU)")
+    gr.Markdown("Upload an image to extract text using OCR or generate text using the vLLM integration.")
 
-    with gr.Row():
-        with gr.Column():
-            prompt_input = gr.Textbox(
-                label="Prompt",
-                placeholder="Enter your prompt here...",
-                lines=3,
-            )
-            max_tokens = gr.Slider(
-                label="Max Tokens",
-                minimum=10,
-                maximum=500,
-                value=100,
-                step=10,
-            )
-            temperature = gr.Slider(
-                label="Temperature",
-                minimum=0.1,
-                maximum=1.0,
-                value=0.7,
-                step=0.1,
-            )
-            top_p = gr.Slider(
-                label="Top P",
-                minimum=0.1,
-                maximum=1.0,
-                value=0.9,
-                step=0.1,
-            )
-            submit_button = gr.Button("Generate")
+    with gr.Tab("Text Generation"):
+        with gr.Row():
+            with gr.Column():
+                prompt_input = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Enter your prompt here...",
+                    lines=3,
+                )
+                max_tokens = gr.Slider(
+                    label="Max Tokens",
+                    minimum=10,
+                    maximum=500,
+                    value=100,
+                    step=10,
+                )
+                temperature = gr.Slider(
+                    label="Temperature",
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.7,
+                    step=0.1,
+                )
+                top_p = gr.Slider(
+                    label="Top P",
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.1,
+                )
+                submit_button = gr.Button("Generate")
+            
+            with gr.Column():
+                output_text = gr.Textbox(
+                    label="Generated Text",
+                    lines=10,
+                    interactive=False,
+                )
         
-        with gr.Column():
-            output_text = gr.Textbox(
-                label="Generated Text",
-                lines=10,
-                interactive=False,
-            )
-    
-    submit_button.click(
-        generate_response,
-        inputs=[prompt_input, max_tokens, temperature, top_p],
-        outputs=output_text,
-    )
+        submit_button.click(
+            generate_response,
+            inputs=[prompt_input, max_tokens, temperature, top_p],
+            outputs=output_text,
+        )
+
+    with gr.Tab("OCR"):
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(
+                    label="Upload Image",
+                    type="file",
+                    image_mode="RGB",
+                )
+                ocr_submit_button = gr.Button("Extract Text")
+            
+            with gr.Column():
+                ocr_output = gr.Textbox(
+                    label="Extracted Text",
+                    lines=10,
+                    interactive=False,
+                )
+        
+        ocr_submit_button.click(
+            ocr_image,
+            inputs=[image_input],
+            outputs=ocr_output,
+        )
 
 # Launch the app
 demo.launch()