fixed few HF warnings

app.py

@@ -2,20 +2,20 @@ import os
 import base64
 from io import BytesIO
 import warnings
+import time  # For timing
 
 import torch
 from PIL import Image
 import gradio as gr
 from transformers import AutoProcessor, AutoModelForVision2Seq
 
-# Suppress warnings at startup
+# Suppress ALL startup noise BEFORE any imports
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["TRANSFORMERS_VERBOSITY"] = "error"
-warnings.filterwarnings("ignore", category=FutureWarning)
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+warnings.filterwarnings("ignore")
 
-# IMPORTANT: Load processor+model from the olmOCR checkpoint itself
 MODEL_ID = "allenai/olmOCR-2-7B-1025"
-
 processor = None
 model = None
 
@@ -25,17 +25,15 @@ def load_model():
     if processor is not None and model is not None:
         return
 
-    # trust_remote_code is often required for VLM checkpoints
     processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-    # T4: use fp16 + device_map auto to avoid OOM
     model = AutoModelForVision2Seq.from_pretrained(
         MODEL_ID,
-        torch_dtype=torch.float16,
+        dtype=torch.float16,  # Fixed deprecation
         device_map="auto",
         low_cpu_mem_usage=True,
         trust_remote_code=True,
     ).eval()
+    print("✅ Model loaded successfully!")
 
 
 def _resize_max_side(img: Image.Image, max_side: int = 896) -> Image.Image:
@@ -48,7 +46,6 @@ def _resize_max_side(img: Image.Image, max_side: int = 896) -> Image.Image:
 
 
 def build_prompt(width: int, height: int) -> str:
-    # Keep it short + strict to reduce hallucinations
     return (
         "Extract all readable text from this page image.\n"
         "Return ONLY the extracted text (no explanations, no markdown).\n"
@@ -59,17 +56,17 @@ def build_prompt(width: int, height: int) -> str:
     )
 
 
-def ocr_image(img: Image.Image) -> str:
+def ocr_image(img: Image.Image) -> tuple[str, str]:
     if img is None:
-        return "No image uploaded."
+        return "No image uploaded.", "0.0s"
 
+    start_time = time.perf_counter()  # High-precision timer
     load_model()
 
     img = img.convert("RGB")
     img = _resize_max_side(img, max_side=896)
     w, h = img.size
 
-    # Base64 image_url message (works with many Qwen-style chat templates)
     buf = BytesIO()
     img.save(buf, format="PNG")
     image_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
@@ -99,27 +96,31 @@ def ocr_image(img: Image.Image) -> str:
         return_tensors="pt",
     )
     
-    # FIX: Move inputs to model's device (eliminates the warning)
+    # Move inputs to model device
     inputs = {k: v.to(model.device) if torch.is_tensor(v) else v for k, v in inputs.items()}
 
     with torch.inference_mode():
         output_ids = model.generate(
             **inputs,
             max_new_tokens=512,
-            do_sample=False,  # deterministic OCR
+            do_sample=False,
         )
 
-    # Remove prompt tokens, keep only generated text
     prompt_len = inputs["input_ids"].shape[1]
     gen_ids = output_ids[:, prompt_len:]
     text_out = processor.tokenizer.batch_decode(gen_ids, skip_special_tokens=True)
-    return text_out[0].strip() if text_out else "No text extracted."
+    result = text_out[0].strip() if text_out else "No text extracted."
+
+    elapsed = time.perf_counter() - start_time
+    timing = f"{elapsed:.2f}s"
+    
+    return result, timing
 
 
 with gr.Blocks(title="BookReader OCR API (olmOCR2)") as demo:
     gr.Markdown(
         "# BookReader OCR API (olmOCR2)\n"
-        "Upload an image → get extracted text.\n\n"
+        "Upload an image → get extracted text + timing.\n\n"
         "**API endpoint:** `/ocr`"
     )
 
@@ -128,12 +129,13 @@ with gr.Blocks(title="BookReader OCR API (olmOCR2)") as demo:
             image_input = gr.Image(type="pil", label="Upload image")
             run_btn = gr.Button("Run OCR", variant="primary")
         with gr.Column():
-            output = gr.Textbox(label="Extracted text", lines=20)
+            output = gr.Textbox(label="Extracted text", lines=15)
+            timing = gr.Textbox(label="Generation time", interactive=False)
 
     run_btn.click(
         fn=ocr_image,
         inputs=[image_input],
-        outputs=[output],
+        outputs=[output, timing],
         api_name="/ocr",
     )