Update app.py

app.py

@@ -1,36 +1,25 @@
-# =========================================================
-# HfFolder shim MUST be first executable code
-# =========================================================
+# ─── 1. Runtime dependency fix (MUST be first) ───
+import bootstrap  # noqa: F401
+
+# ─── 2. HfFolder shim for old Gradio ───
 import huggingface_hub
 if not hasattr(huggingface_hub, "HfFolder"):
     class _HfFolder:
         @staticmethod
         def get_token(): return huggingface_hub.get_token()
         @staticmethod
-        def save_token(token): huggingface_hub.login(token)
+        def save_token(t): huggingface_hub.login(t)
         @staticmethod
         def delete_token(): huggingface_hub.logout()
     huggingface_hub.HfFolder = _HfFolder
 
-# =========================================================
-# All other imports
-# =========================================================
-import os
-import re
-import csv
-import json
-import sqlite3
-import tempfile
-import logging
-import asyncio
-import uuid
-import threading
-import httpx
-import pymupdf
+# ─── 3. spaces BEFORE torch ───
 import spaces
+import torch
 import gradio as gr
+import os, re, csv, json, sqlite3, tempfile, logging, asyncio, uuid, httpx
 from PIL import Image
-from gradio_client import Client, handle_file as gradio_handle_file
+from transformers import AutoModel, AutoTokenizer
 from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
 from telegram.ext import (
     Application, CommandHandler, MessageHandler, filters,
@@ -41,319 +30,204 @@ from fastapi import Request
 from fastapi.responses import JSONResponse
 from starlette.routing import Route as StarletteRoute
 
-# --- gradio_client fetch timeout ---
-import gradio_client.client as gradio_client_module
-_original_httpx_get = httpx.get
-def _httpx_get_with_timeout(*args, **kwargs):
-    if "timeout" not in kwargs:
-        kwargs["timeout"] = httpx.Timeout(120.0, connect=10.0)
-    return _original_httpx_get(*args, **kwargs)
-gradio_client_module.httpx.get = _httpx_get_with_timeout
-
-# =========================================================
-# Configuration
-# =========================================================
-TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
-HF_SPACE_ID = "khang119966/DeepSeek-OCR-DEMO"
-DB_PATH = "ocr_records.db"
-WEBHOOK_URL = os.getenv("WEBHOOK_URL", "").rstrip("/")
-
-if not TELEGRAM_BOT_TOKEN:
-    raise RuntimeError("Missing TELEGRAM_BOT_TOKEN")
-if not WEBHOOK_URL:
-    raise RuntimeError("Missing WEBHOOK_URL (e.g. https://your-space.hf.space)")
-
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-# =========================================================
-# Gradio Client
-# =========================================================
-try:
-    gradio_client = Client(HF_SPACE_ID, httpx_kwargs={"timeout": httpx.Timeout(120.0, connect=10.0)})
-    logger.info(f"Connected to OCR Space: {HF_SPACE_ID}")
-except Exception as e:
-    logger.error(f"Failed to init Gradio Client: {e}")
-    raise e
-
-# =========================================================
-# Database
-# =========================================================
+# ─── 4. Config ───
+TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
+WEBHOOK_URL = os.getenv("WEBHOOK_URL", "").rstrip("/")
+DB_PATH = "ocr_records.db"
+MODEL_ID = "strangervisionhf/deepseek-ocr-transformers-v4.57.6"
+GPU_DURATION = 30  # seconds reserved per OCR call
+
+if not TELEGRAM_BOT_TOKEN or not WEBHOOK_URL:
+    raise RuntimeError("TELEGRAM_BOT_TOKEN and WEBHOOK_URL are required")
+
+# ─── 5. Load model on CPU at module level ───
+logger.info("Loading DeepSeek-OCR on CPU...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    MODEL_ID,
+    _attn_implementation="eager",
+    trust_remote_code=True,
+    use_safetensors=True,
+).eval()
+logger.info("✅ Model loaded (CPU).")
+
+# ─── 6. DB + parsing (unchanged from your current app) ───
 def init_db():
     conn = sqlite3.connect(DB_PATH)
-    cur = conn.cursor()
-    cur.execute("""
-        CREATE TABLE IF NOT EXISTS voters (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            voter_id TEXT, name TEXT, father_name TEXT, husband_name TEXT,
-            house_no TEXT, age INTEGER, gender TEXT, raw_json TEXT
-        )
-    """)
-    conn.commit()
-    conn.close()
-
-# =========================================================
-# Parsing
-# =========================================================
-def parse_markdown_to_records(raw_markdown: str) -> list[dict]:
-    table_match = re.search(r'<table>.*?</table>', raw_markdown, re.DOTALL)
-    if not table_match:
-        return []
-    html_content = table_match.group(0)
-    cells = re.findall(r'<td>(.*?)</td>', html_content, re.DOTALL)
-    id_pattern = re.compile(r'(UP/\d+/\d+/\d+|[A-Z]{3}\d{7,})')
-    voter_ids = [c for c in cells if id_pattern.fullmatch(c.strip())]
-    detail_cells = [c for c in cells if "नाम:" in c]
-    records = []
-    for i in range(min(len(voter_ids), len(detail_cells))):
-        voter_id = voter_ids[i].strip()
-        text = re.sub(r'<br\s*/?>', '\n', detail_cells[i])
+    conn.execute("""CREATE TABLE IF NOT EXISTS voters (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        voter_id TEXT, name TEXT, father_name TEXT, husband_name TEXT,
+        house_no TEXT, age INTEGER, gender TEXT, raw_json TEXT)""")
+    conn.commit(); conn.close()
+
+def parse_markdown(raw: str) -> list[dict]:
+    m = re.search(r'<table>.*?</table>', raw, re.DOTALL)
+    if not m: return []
+    cells = re.findall(r'<td>(.*?)</td>', m.group(0), re.DOTALL)
+    ids = [c for c in cells if re.fullmatch(r'(UP/\d+/\d+/\d+|[A-Z]{3}\d{7,})', c.strip())]
+    details = [c for c in cells if "नाम:" in c]
+    out = []
+    for i in range(min(len(ids), len(details))):
+        text = re.sub(r'<br\s*/?>', '\n', details[i])
         text = re.sub(r'<[^>]+>', '', text)
-        nm = re.search(r'नाम:\s*([^\n]+)', text)
-        fm = re.search(r'पिता का नाम:\s*([^\n]+)', text)
-        hm = re.search(r'पति का नाम:\s*([^\n]+)', text)
-        hm_no = re.search(r'मकान संख्या:\s*([^\n]+)', text)
-        ag = re.search(r'आयु:\s*(\d+)', text)
-        gm = re.search(r'लिंग:\s*([^\n]+)', text)
-        records.append({
-            "voter_id": voter_id,
-            "name": nm.group(1).strip() if nm else "",
-            "father_name": fm.group(1).strip() if fm else "",
-            "husband_name": hm.group(1).strip() if hm else "",
-            "house_no": hm_no.group(1).strip() if hm_no else "",
-            "age": int(ag.group(1)) if ag else 0,
-            "gender": gm.group(1).strip() if gm else "",
+        def f(p): 
+            r = re.search(p, text); return r.group(1).strip() if r else ""
+        age = re.search(r'आयु:\s*(\d+)', text)
+        out.append({
+            "voter_id": ids[i].strip(),
+            "name": f(r'नाम:\s*([^\n]+)'),
+            "father_name": f(r'पिता का नाम:\s*([^\n]+)'),
+            "husband_name": f(r'पति का नाम:\s*([^\n]+)'),
+            "house_no": f(r'मकान संख्या:\s*([^\n]+)'),
+            "age": int(age.group(1)) if age else 0,
+            "gender": f(r'लिंग:\s*([^\n]+)'),
         })
-    return records
+    return out
 
-# =========================================================
-# OCR call
-# =========================================================
-def call_space_ocr(image_path: str) -> str:
-    clean = os.path.join(tempfile.gettempdir(), f"ocr_input_{uuid.uuid4().hex}.jpg")
+# ─── 7. GPU OCR call ───
+@spaces.GPU(duration=GPU_DURATION)
+def run_ocr(image_path: str) -> str:
+    """Runs on GPU. Model moves to VRAM, then back to CPU to free quota."""
+    model.cuda().to(torch.bfloat16)
     try:
-        with Image.open(image_path) as img:
-            img.convert("RGB").save(clean, "JPEG")
-        result = gradio_client.predict(
-            image=gradio_handle_file(clean),
-            model_size="Gundam (Recommended)",
-            task_type="📄 Convert to Markdown",
-            ref_text=" ",
-            api_name="/process_ocr_task",
-        )
-        return result[0] if isinstance(result, tuple) else str(result)
-    except Exception as exc:
-        logger.exception(f"OCR failed: {image_path}")
-        raise exc
+        with tempfile.TemporaryDirectory() as out:
+            result = model.infer(
+                tokenizer,
+                prompt="<image>\n<|grounding|>Convert the document to markdown.",
+                image_file=image_path,
+                output_path=out,
+                base_size=1024, image_size=640, crop_mode=True,
+                save_results=False, test_compress=True, eval_mode=True,
+            )
+        return result
     finally:
-        if os.path.exists(clean):
-            try: os.remove(clean)
-            except OSError: pass
+        model.cpu()  # free VRAM immediately
 
-# =========================================================
-# Progress bar
-# =========================================================
-async def animate_progress(status_msg, stop_event: asyncio.Event):
-    bar_len = 10
-    i = 0
-    while not stop_event.is_set():
-        bar = "█" * i + "░" * (bar_len - i)
-        try:
-            await status_msg.edit_text(f"⚙️ OCR in progress...\n`[{bar}]`", parse_mode="Markdown")
-        except BadRequest:
-            pass
-        i = (i + 1) % (bar_len + 1)
-        await asyncio.sleep(2)
+# ─── 8. Telegram handlers ───
+async def start(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
+    await update.message.reply_text("👋 Send an image — I'll OCR it locally and let you verify before saving.")
 
-# =========================================================
-# Telegram handlers
-# =========================================================
-async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    await update.message.reply_text("👋 Send me an Image! I'll extract voter data and let you verify before saving.")
-
-async def handle_file(update: Update, context: ContextTypes.DEFAULT_TYPE):
+async def on_image(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
     msg = update.message
     if msg.document:
         if not msg.document.mime_type.startswith("image/"):
-            await msg.reply_text("❌ Please send an Image.")
-            return
-        tg_file = await context.bot.get_file(msg.document)
-        original_name = msg.document.file_name or "upload.jpg"
+            await msg.reply_text("❌ Please send an image."); return
+        f = await ctx.bot.get_file(msg.document)
+        name = msg.document.file_name or "img.jpg"
     elif msg.photo:
-        tg_file = await context.bot.get_file(msg.photo[-1].file_id)
-        original_name = f"photo_{msg.message_id}.jpg"
-    else:
-        return
+        f = await ctx.bot.get_file(msg.photo[-1].file_id)
+        name = f"photo_{msg.message_id}.jpg"
+    else: return
 
-    safe_name = re.sub(r"[^\w.\-]", "_", original_name)
-    download_path = os.path.join(tempfile.gettempdir(), safe_name)
-    await tg_file.download_to_drive(download_path)
-    status_msg = await msg.reply_text("📥 File received. Preparing image...")
+    path = os.path.join(tempfile.gettempdir(), re.sub(r"[^\w.\-]", "_", name))
+    await f.download_to_drive(path)
+    status = await msg.reply_text("⚙️ Running OCR on GPU… (cold start may take 60s)")
 
     try:
-        stop_event = asyncio.Event()
-        progress_task = asyncio.create_task(animate_progress(status_msg, stop_event))
-        raw_text = await asyncio.to_thread(call_space_ocr, download_path)
-        stop_event.set()
-        await progress_task
-        await status_msg.edit_text("✅ OCR Complete! Parsing data...", parse_mode="Markdown")
-
-        records = parse_markdown_to_records(raw_text)
+        raw = await asyncio.to_thread(run_ocr, path)
+        records = parse_markdown(raw)
         if not records:
-            await msg.reply_text("⚠️ No structured records found. Raw text:\n\n" + raw_text[:3000])
-            return
+            await msg.reply_text("⚠️ No table found. Raw:\n\n" + raw[:3000]); return
 
-        csv_path = os.path.join(tempfile.gettempdir(), f"ocr_results_{uuid.uuid4().hex}.csv")
-        with open(csv_path, "w", newline="", encoding="utf-8-sig") as f:
-            writer = csv.DictWriter(f, fieldnames=["voter_id", "name", "father_name", "husband_name", "house_no", "age", "gender"])
-            writer.writeheader()
-            writer.writerows(records)
+        csv_path = os.path.join(tempfile.gettempdir(), f"out_{uuid.uuid4().hex}.csv")
+        with open(csv_path, "w", newline="", encoding="utf-8-sig") as fh:
+            w = csv.DictWriter(fh, fieldnames=["voter_id","name","father_name","husband_name","house_no","age","gender"])
+            w.writeheader(); w.writerows(records)
 
-        context.user_data['temp_csv_path'] = csv_path
-        context.user_data['temp_records'] = records
-
-        keyboard = [[
-            InlineKeyboardButton("✅ Save to DB", callback_data="save_db"),
-            InlineKeyboardButton("❌ Cancel", callback_data="cancel"),
-        ]]
+        ctx.user_data['csv'] = csv_path
+        ctx.user_data['records'] = records
+        kb = [[InlineKeyboardButton("✅ Save to DB", callback_data="save"),
+               InlineKeyboardButton("❌ Cancel", callback_data="cancel")]]
         await msg.reply_document(
-            document=open(csv_path, "rb"),
-            filename="extracted_voters.csv",
-            caption=f"📊 Extracted **{len(records)}** records.\nPlease verify. Click **Save to DB** if correct.",
-            parse_mode="Markdown",
-            reply_markup=InlineKeyboardMarkup(keyboard),
+            open(csv_path, "rb"), filename="extracted_voters.csv",
+            caption=f"📊 Extracted **{len(records)}** records. Verify, then Save.",
+            parse_mode="Markdown", reply_markup=InlineKeyboardMarkup(kb),
         )
-    except Exception as exc:
-        if 'stop_event' in locals():
-            stop_event.set()
-            await progress_task
-        await status_msg.edit_text(f"❌ Error: {exc}", parse_mode="Markdown")
+    except Exception as e:
+        logger.exception("OCR pipeline failed")
+        await status.edit_text(f"❌ Error: {e}")
     finally:
-        if os.path.exists(download_path):
-            os.remove(download_path)
-
-async def button_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    query = update.callback_query
-    await query.answer()
-    csv_path = context.user_data.get('temp_csv_path')
-    records = context.user_data.get('temp_records')
-
-    if query.data == "save_db" and records:
-        try:
-            conn = sqlite3.connect(DB_PATH)
-            cur = conn.cursor()
-            for r in records:
-                cur.execute("""
-                    INSERT INTO voters (voter_id, name, father_name, husband_name, house_no, age, gender, raw_json)
-                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
-                """, (r["voter_id"], r["name"], r["father_name"], r["husband_name"],
-                      r["house_no"], r["age"], r["gender"], json.dumps(r, ensure_ascii=False)))
-            conn.commit()
-            conn.close()
-            await query.edit_message_caption(caption=f"✅ Success! {len(records)} records saved to DB.")
-        except Exception as e:
-            await query.edit_message_caption(caption=f"❌ Database Error: {e}")
-    elif query.data == "cancel":
-        await query.edit_message_caption(caption="❌ Cancelled. Data not saved.")
-
-    if csv_path and os.path.exists(csv_path):
-        os.remove(csv_path)
-    context.user_data.clear()
-
-# =========================================================
-# Telegram Application (webhook mode)
-# =========================================================
-telegram_app = (
-    Application.builder()
-    .token(TELEGRAM_BOT_TOKEN)
-    .updater(None)
-    .build()
-)
+        if os.path.exists(path): os.remove(path)
+
+async def on_button(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
+    q = update.callback_query; await q.answer()
+    csv_path = ctx.user_data.get('csv')
+    records = ctx.user_data.get('records')
+    if q.data == "save" and records:
+        conn = sqlite3.connect(DB_PATH)
+        for r in records:
+            conn.execute("""INSERT INTO voters
+                (voter_id,name,father_name,husband_name,house_no,age,gender,raw_json)
+                VALUES (?,?,?,?,?,?,?,?)""",
+                (r["voter_id"], r["name"], r["father_name"], r["husband_name"],
+                 r["house_no"], r["age"], r["gender"], json.dumps(r, ensure_ascii=False)))
+        conn.commit(); conn.close()
+        await q.edit_message_caption(caption=f"✅ Saved {len(records)} records to DB.")
+    elif q.data == "cancel":
+        await q.edit_message_caption(caption="❌ Cancelled.")
+    if csv_path and os.path.exists(csv_path): os.remove(csv_path)
+    ctx.user_data.clear()
+
+# ─── 9. Telegram app (webhook mode) ───
+telegram_app = Application.builder().token(TELEGRAM_BOT_TOKEN).updater(None).build()
 telegram_app.add_handler(CommandHandler("start", start))
-telegram_app.add_handler(MessageHandler((filters.Document.ALL | filters.PHOTO) & ~filters.COMMAND, handle_file))
-telegram_app.add_handler(CallbackQueryHandler(button_callback))
-
-# =========================================================
-# Dedicated event loop for Telegram (separate from Uvicorn)
-# =========================================================
-_telegram_loop = asyncio.new_event_loop()
-_telegram_ready = threading.Event()
+telegram_app.add_handler(MessageHandler((filters.Document.ALL | filters.PHOTO) & ~filters.COMMAND, on_image))
+telegram_app.add_handler(CallbackQueryHandler(on_button))
 
-async def _init_telegram():
-    init_db()
-    await telegram_app.initialize()
-    await telegram_app.start()
-    webhook_url = f"{WEBHOOK_URL}/webhook/{TELEGRAM_BOT_TOKEN}"
-    await telegram_app.bot.set_webhook(url=webhook_url, drop_pending_updates=True)
-    logger.info(f"✅ Webhook set: {webhook_url}")
-
-def _run_telegram_loop():
-    asyncio.set_event_loop(_telegram_loop)
-    _telegram_loop.run_until_complete(_init_telegram())
-    _telegram_ready.set()
-    _telegram_loop.run_forever()
-
-_thread = threading.Thread(target=_run_telegram_loop, daemon=True)
-_thread.start()
-
-# Wait (up to 60s) for Telegram to be fully initialized before accepting webhooks
-if not _telegram_ready.wait(timeout=60):
-    logger.error("Telegram init did not complete within 60s")
-else:
-    logger.info("Telegram loop is ready.")
-
-# =========================================================
-# Webhook & health handlers
-# =========================================================
-async def telegram_webhook(request: Request):
+# ─── 10. Dedicated Telegram loop ───
+import threading
+_loop = asyncio.new_event_loop()
+_ready = threading.Event()
+
+def _run():
+    asyncio.set_event_loop(_loop)
+    async def _init():
+        init_db()
+        await telegram_app.initialize()
+        await telegram_app.start()
+        await telegram_app.bot.set_webhook(
+            url=f"{WEBHOOK_URL}/webhook/{TELEGRAM_BOT_TOKEN}",
+            drop_pending_updates=True)
+        logger.info(f"✅ Webhook set: {WEBHOOK_URL}/webhook/{TELEGRAM_BOT_TOKEN}")
+    _loop.run_until_complete(_init())
+    _ready.set()
+    _loop.run_forever()
+
+threading.Thread(target=_run, daemon=True).start()
+_ready.wait(timeout=60)
+
+# ─── 11. Webhook + health routes ───
+async def telegram_webhook(req: Request):
     try:
-        data = await request.json()
+        data = await req.json()
         update = Update.de_json(data, telegram_app.bot)
-        # Bridge: schedule process_update on Telegram's own loop
-        future = asyncio.run_coroutine_threadsafe(
-            telegram_app.process_update(update),
-            _telegram_loop,
-        )
-        await asyncio.wait_for(asyncio.wrap_future(future), timeout=120.0)
+        fut = asyncio.run_coroutine_threadsafe(telegram_app.process_update(update), _loop)
+        await asyncio.wait_for(asyncio.wrap_future(fut), timeout=120)
     except Exception as e:
-        logger.exception(f"Webhook handler error: {e}")
+        logger.exception(f"Webhook error: {e}")
     return JSONResponse({"ok": True})
 
-async def health_handler(request: Request):
+async def health(req: Request):
     return JSONResponse({"status": "ok"})
 
-# =========================================================
-# ZeroGPU hidden button
-# =========================================================
-@spaces.GPU(duration=10)
-def _zerogpu_register():
-    return "gpu-ready"
+# ─── 12. Hidden GPU button (ZeroGPU scanner needs it) ───
+@spaces.GPU(duration=1)
+def _probe(): return "ok"
 
-with gr.Blocks(title="Telegram OCR Bot") as demo:
-    gr.Markdown("# ✅ Telegram OCR Bot — keep this Space awake by pinging `/health`")
-    _btn = gr.Button("init-gpu", visible=False)
-    _btn.click(fn=_zerogpu_register, inputs=None, outputs=None)
+with gr.Blocks() as demo:
+    gr.Markdown("# DeepSeek-OCR Telegram Bot — ping `/health` to keep alive")
+    _b = gr.Button("init-gpu", visible=False)
+    _b.click(fn=_probe)
 
-# =========================================================
-# Main
-# =========================================================
+# ─── 13. Launch ───
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        ssr_mode=False,
-        prevent_thread_lock=True,
-    )
-
-    # Inject routes at index 0 (beat Gradio's catch-all)
-    live_app = demo.app
-    live_app.router.routes.insert(
-        0, StarletteRoute(f"/webhook/{TELEGRAM_BOT_TOKEN}", telegram_webhook, methods=["POST"])
-    )
-    live_app.router.routes.insert(
-        0, StarletteRoute("/health", health_handler, methods=["GET"])
-    )
-    live_app.build_middleware_stack()
-
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False,
+                ssr_mode=False, prevent_thread_lock=True)
+    live = demo.app
+    live.router.routes.insert(0, StarletteRoute(f"/webhook/{TELEGRAM_BOT_TOKEN}", telegram_webhook, methods=["POST"]))
+    live.router.routes.insert(0, StarletteRoute("/health", health, methods=["GET"]))
+    live.build_middleware_stack()
     demo.block_thread()