Update app.py

app.py

@@ -1,42 +1,21 @@
 import torch
+import gradio as gr
 from transformers import AutoProcessor, AutoModelForVision2Seq
 from PIL import Image
-from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import JSONResponse
-import io
-import transformers
-print(f"Transformers version: {transformers.__version__}")
 
-app = FastAPI()
-
-# SmolVLM - lightweight, works on CPU/low VRAM, same team as Idefics3
 model_id = "HuggingFaceTB/SmolVLM-Instruct"
 
-print("Loading processor...")
-processor = AutoProcessor.from_pretrained(model_id)
-
 print("Loading model...")
+processor = AutoProcessor.from_pretrained(model_id)
 model = AutoModelForVision2Seq.from_pretrained(
     model_id,
-    torch_dtype=torch.float32,   # float32 for CPU
+    torch_dtype=torch.float32,
     device_map="auto"
 )
 model.eval()
 print("Model ready!")
 
-
-@app.get("/")
-def root():
-    return {"status": "running", "model": model_id}
-
-
-@app.post("/extract")
-async def extract_text(file: UploadFile = File(...)):
-    # Read uploaded image
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
-
-    # Prompt for NEET question extraction
+def extract_text(image: Image.Image) -> str:
     messages = [
         {
             "role": "user",
@@ -44,7 +23,7 @@ async def extract_text(file: UploadFile = File(...)):
                 {"type": "image"},
                 {"type": "text", "text": (
                     "Extract all text from this image exactly as it appears. "
-                    "Preserve question numbers, options (A, B, C, D), tables, "
+                    "Preserve question numbers, options A B C D, tables, "
                     "and any mathematical or chemical expressions. "
                     "Format clearly."
                 )}
@@ -53,12 +32,7 @@ async def extract_text(file: UploadFile = File(...)):
     ]
 
     prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-
-    inputs = processor(
-        text=prompt,
-        images=[image],
-        return_tensors="pt"
-    )
+    inputs = processor(text=prompt, images=[image], return_tensors="pt")
 
     with torch.no_grad():
         outputs = model.generate(
@@ -67,8 +41,16 @@ async def extract_text(file: UploadFile = File(...)):
             do_sample=False
         )
 
-    # Decode only the generated part
     generated = outputs[0][inputs["input_ids"].shape[1]:]
-    result = processor.decode(generated, skip_special_tokens=True)
+    return processor.decode(generated, skip_special_tokens=True)
+
+
+demo = gr.Interface(
+    fn=extract_text,
+    inputs=gr.Image(type="pil", label="Upload NEET Question Image"),
+    outputs=gr.Textbox(label="Extracted Text", lines=20),
+    title="NEET Question Extractor",
+    description="Upload a scanned NEET question paper image to extract text"
+)
 
-    return JSONResponse({"extracted_text": result})
+demo.launch()