app.py

"""
Resume OCR — Local Qwen2.5-VL edition
- Gradio UI at /
- FastAPI REST endpoint at /extract-resume
- FastAPI docs at /docs
"""

import os
import re
import json
import time
import torch
import gradio as gr
import fitz  # pymupdf
import uvicorn
from contextlib import asynccontextmanager
from io import BytesIO
from PIL import Image
from typing import List
from pydantic import BaseModel
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

# ── Cache dir ──────────────────────────────────────────────────────────────────
def get_cache_dir():
    for candidate in ["/data/hf_cache", "/tmp/hf_cache"]:
        try:
            os.makedirs(candidate, exist_ok=True)
            test = os.path.join(candidate, ".write_test")
            with open(test, "w") as f:
                f.write("ok")
            os.remove(test)
            return candidate
        except (PermissionError, OSError):
            continue
    raise RuntimeError("No writable cache directory found.")

CACHE_DIR = get_cache_dir()
os.environ["HF_HOME"] = CACHE_DIR
os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR

# ── Model config ───────────────────────────────────────────────────────────────
MODEL_ID        = "Qwen/Qwen2.5-VL-3B-Instruct"
CPU_THREADS     = 4
DEFAULT_MIN_PX  = 256
DEFAULT_MAX_PX  = 512
DEFAULT_MAX_TOK = 2048
PDF_MAX_PAGES   = 10
ALLOWED_EXTENSIONS = (".jpg", ".jpeg", ".png", ".webp", ".pdf")

torch.set_num_threads(CPU_THREADS)

# Global model references — populated in lifespan
model     = None
processor = None

# ── Lifespan: load model once on startup ──────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    global model, processor
    print(f"\n{'='*60}")
    print(f"  Loading {MODEL_ID}")
    print(f"  Device: CPU   Threads: {CPU_THREADS}   dtype: float32")
    print(f"  Cache : {CACHE_DIR}")
    print(f"{'='*60}\n")
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,
        device_map="cpu",
        cache_dir=CACHE_DIR,
    )
    model.eval()
    processor = AutoProcessor.from_pretrained(
        MODEL_ID,
        min_pixels=DEFAULT_MIN_PX * 28 * 28,
        max_pixels=DEFAULT_MAX_PX * 28 * 28,
        cache_dir=CACHE_DIR,
    )
    print("✓ Model ready\n")
    yield
    # cleanup (none needed)

# ── RESUME_SYSTEM_PROMPT ───────────────────────────────────────────────────────
RESUME_SYSTEM_PROMPT = """You are a resume data extraction assistant.
You will receive a resume/CV image (or images for multi-page PDFs).
Parse it carefully and return ONLY a valid JSON object.
No markdown fences, no explanation, no preamble — just the raw JSON object.
JSON schema (return exactly this structure):
{
  "name": "full name of the candidate (string)",
  "email": "email address (string)",
  "phone": "phone number as a string",
  "address": "full address or location (string)",
  "marital_status": "marital status e.g. Single, Married, Divorced (string)",
  "gender": "gender e.g. Male, Female, Other (string)",
  "location": "city, state or country (string)",
  "linkedin": "LinkedIn profile URL or username (string)",
  "summary": "professional summary or objective, 1-3 sentences (string)",
  "skills": ["skill1", "skill2", "skill3"],
  "languages": ["language1", "language2"],
  "years_of_experience": "total years of experience as a number or string (string)",
  "experience": [
    {
      "title": "job title (string)",
      "company": "company name (string)",
      "duration": "employment period e.g. Jan 2020 - Mar 2023 (string)",
      "description": "brief summary of responsibilities and achievements (string)"
    }
  ],
  "education": [
    {
      "degree": "degree name e.g. B.Tech Computer Science (string)",
      "institution": "university or college name (string)",
      "year": "graduation year or period (string)",
      "grade": "GPA, percentage, or grade if mentioned (string)"
    }
  ],
  "certifications": ["certification1", "certification2"],
  "projects": [
    {
      "name": "project name (string)",
      "description": "brief description (string)",
      "technologies": "technologies used (string)"
    }
  ]
}
Rules:
- name: the candidate's full name, usually at the top of the resume in large text
- email: look for @ symbol; return empty string if not found
- phone: may include or may not include country code; must be 10 digits; return as string
- address: look for postal address, street address, or location information
- marital_status: look for sections or lines mentioning Single, Married, Divorced, etc.
- gender: look for gender information if mentioned; return empty string if not found
- location: city and country or state; do not include full street address
- linkedin: look for linkedin.com URL or "LinkedIn:" label
- summary: look for sections labeled "Summary", "Objective", "About", or "Profile"
- skills: extract all technical and soft skills as a flat list of strings
- languages: list all spoken/written languages mentioned
- years_of_experience: calculate from work experience dates or look for explicit mention
- experience: list all jobs in order from most recent; include internships
- education: list all degrees/diplomas; include ongoing education
- certifications: list all certifications, courses, or training
- projects: list notable projects; include personal, academic, or professional
- If a field is not found, use "" for strings and [] for arrays
- Do NOT invent or hallucinate any information not present in the image"""

# ── Pydantic models ────────────────────────────────────────────────────────────
class ExperienceItem(BaseModel):
    title: str
    company: str
    duration: str
    description: str

class EducationItem(BaseModel):
    degree: str
    institution: str
    year: str
    grade: str

class ProjectItem(BaseModel):
    name: str
    description: str
    technologies: str

class ResumeData(BaseModel):
    name: str
    email: str
    phone: str
    address: str
    marital_status: str
    gender: str
    location: str
    linkedin: str
    summary: str
    skills: List[str]
    experience: List[ExperienceItem]
    education: List[EducationItem]
    certifications: List[str]
    languages: List[str]
    years_of_experience: str
    projects: List[ProjectItem]

# ── PDF → PIL list ─────────────────────────────────────────────────────────────
def pdf_bytes_to_pil_list(pdf_bytes: bytes, dpi: int = 200) -> list:
    try:
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    except Exception as exc:
        raise ValueError(f"Could not read PDF: {exc}")
    if doc.page_count == 0:
        raise ValueError("PDF has no pages.")
    if doc.page_count > PDF_MAX_PAGES:
        raise ValueError(f"PDF has {doc.page_count} pages; max is {PDF_MAX_PAGES}.")
    images = []
    for page in doc:
        pix = page.get_pixmap(dpi=dpi)
        images.append(Image.frombytes("RGB", [pix.width, pix.height], pix.samples))
    doc.close()
    return images

# ── JSON cleaning ──────────────────────────────────────────────────────────────
def clean_and_parse_json(raw_text: str) -> dict:
    cleaned = re.sub(r"```json\s*", "", raw_text, flags=re.IGNORECASE)
    cleaned = re.sub(r"```\s*", "", cleaned).strip()
    try:
        parsed = json.loads(cleaned)
    except json.JSONDecodeError:
        match = re.search(r"\{[\s\S]*\}", cleaned)
        if not match:
            raise ValueError(f"No JSON found. Preview: {raw_text[:400]}")
        parsed = json.loads(match.group(0))
    if not isinstance(parsed, dict):
        raise ValueError(f"Not a JSON object. Got: {type(parsed).__name__}")
    return parsed

# ── Build ResumeData ───────────────────────────────────────────────────────────
def build_resume_data(parsed: dict) -> ResumeData:
    def safe_str(key, max_len=200):
        return str(parsed.get(key, "")).strip()[:max_len]
    def safe_list(key):
        val = parsed.get(key, [])
        return [str(x).strip() for x in val if x] if isinstance(val, list) else []
    return ResumeData(
        name=safe_str("name", 100),
        email=safe_str("email", 100),
        phone=safe_str("phone", 20),
        address=safe_str("address", 200),
        marital_status=safe_str("marital_status", 50),
        gender=safe_str("gender", 50),
        location=safe_str("location", 100),
        linkedin=safe_str("linkedin", 150),
        summary=safe_str("summary", 1000),
        skills=safe_list("skills"),
        experience=[
            ExperienceItem(
                title=str(i.get("title","")).strip()[:100],
                company=str(i.get("company","")).strip()[:100],
                duration=str(i.get("duration","")).strip()[:60],
                description=str(i.get("description","")).strip()[:500],
            ) for i in parsed.get("experience",[]) if isinstance(i, dict)
        ],
        education=[
            EducationItem(
                degree=str(i.get("degree","")).strip()[:150],
                institution=str(i.get("institution","")).strip()[:150],
                year=str(i.get("year","")).strip()[:30],
                grade=str(i.get("grade","")).strip()[:30],
            ) for i in parsed.get("education",[]) if isinstance(i, dict)
        ],
        certifications=safe_list("certifications"),
        languages=safe_list("languages"),
        years_of_experience=safe_str("years_of_experience", 50),
        projects=[
            ProjectItem(
                name=str(i.get("name","")).strip()[:100],
                description=str(i.get("description","")).strip()[:500],
                technologies=str(i.get("technologies","")).strip()[:200],
            ) for i in parsed.get("projects",[]) if isinstance(i, dict)
        ],
    )

# ── Qwen inference ─────────────────────────────────────────────────────────────
def run_qwen(images: list, min_px: int, max_px: int, max_new_tokens: int) -> str:
    content = []
    for img in images:
        content.append({
            "type": "image",
            "image": img,
            "min_pixels": min_px * 28 * 28,
            "max_pixels": max_px * 28 * 28,
        })
    content.append({"type": "text", "text": RESUME_SYSTEM_PROMPT})
    messages = [{"role": "user", "content": content}]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(text=[text], images=image_inputs, videos=video_inputs,
                       padding=True, return_tensors="pt")
    t0 = time.time()
    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs, max_new_tokens=max_new_tokens,
            do_sample=False, temperature=None, top_p=None,
        )
    elapsed = time.time() - t0
    trimmed = [out[len(inp):] for inp, out in zip(inputs.input_ids, generated_ids)]
    raw_text = processor.batch_decode(trimmed, skip_special_tokens=True,
                                      clean_up_tokenization_spaces=False)[0]
    print(f"  ⏱ {elapsed:.1f}s | {len(trimmed[0])} tokens | {len(trimmed[0])/elapsed:.2f} tok/s")
    return raw_text

# ── Shared pipeline ────────────────────────────────────────────────────────────
def pipeline_from_bytes(raw_bytes: bytes, filename: str,
                        min_px=DEFAULT_MIN_PX, max_px=DEFAULT_MAX_PX,
                        max_new_tokens=DEFAULT_MAX_TOK) -> ResumeData:
    is_pdf = raw_bytes[:4] == b"%PDF" or filename.lower().endswith(".pdf")
    images = pdf_bytes_to_pil_list(raw_bytes) if is_pdf \
             else [Image.open(BytesIO(raw_bytes)).convert("RGB")]
    raw_text = run_qwen(images, min_px, max_px, max_new_tokens)
    parsed   = clean_and_parse_json(raw_text)
    return build_resume_data(parsed)

# ── FastAPI ────────────────────────────────────────────────────────────────────
api = FastAPI(title="Resume OCR API — Local Qwen", lifespan=lifespan)

api.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], allow_credentials=True,
    allow_methods=["*"], allow_headers=["*"],
)

@api.get("/health")
async def health():
    return {"status": "healthy", "model": MODEL_ID}

@api.post("/extract-resume", response_model=ResumeData)
async def extract_resume_api(file: UploadFile = File(...)):
    filename = (file.filename or "").lower()
    if not filename.endswith(ALLOWED_EXTENSIONS):
        raise HTTPException(415, f"Unsupported file type: {file.content_type}")
    raw_bytes = await file.read()
    try:
        return pipeline_from_bytes(raw_bytes, filename)
    except ValueError as e:
        raise HTTPException(422, str(e))
    except Exception as e:
        raise HTTPException(500, str(e))

# ── Pretty-print ───────────────────────────────────────────────────────────────
def format_resume(d: dict) -> str:
    lines = [f"{'='*60}", f"  {d.get('name','—')}", f"{'='*60}"]
    contact = [d.get(k,"").strip() for k in ("email","phone","location","linkedin") if d.get(k,"").strip()]
    if contact:
        lines.append("  " + "  |  ".join(contact))
    for key, label in (("gender","Gender"),("marital_status","Marital status"),("address","Address")):
        if d.get(key,"").strip():
            lines.append(f"  {label}: {d[key].strip()}")
    if d.get("summary","").strip():
        lines += ["", "SUMMARY", "-------", d["summary"].strip()]
    if d.get("experience"):
        lines += ["", "EXPERIENCE", "----------"]
        for exp in d["experience"]:
            lines.append(
                f"  {exp.get('title','')}"
                + (f" @ {exp['company']}" if exp.get("company") else "")
                + (f"  [{exp['duration']}]" if exp.get("duration") else "")
            )
            for s in (exp.get("description","") or "").split(". "):
                if s.strip(): lines.append(f"    • {s.strip().rstrip('.')}.")
            lines.append("")
    if d.get("education"):
        lines += ["EDUCATION", "---------"]
        for edu in d["education"]:
            lines.append(
                f"  {edu.get('degree','')}"
                + (f" — {edu['institution']}" if edu.get("institution") else "")
                + (f"  ({edu['year']})" if edu.get("year") else "")
                + (f"  [{edu['grade']}]" if edu.get("grade") else "")
            )
        lines.append("")
    skills = [s.strip() for s in d.get("skills",[]) if s.strip()]
    if skills:
        lines += ["SKILLS", "------"]
        row, cur = [], 0
        for s in skills:
            if cur + len(s) + 2 > 70 and row:
                lines.append("  " + ", ".join(row)); row, cur = [], 0
            row.append(s); cur += len(s) + 2
        if row: lines.append("  " + ", ".join(row))
        lines.append("")
    certs = [c.strip() for c in d.get("certifications",[]) if c.strip()]
    if certs:
        lines += ["CERTIFICATIONS", "--------------"] + [f"  • {c}" for c in certs] + [""]
    langs = [l.strip() for l in d.get("languages",[]) if l.strip()]
    if langs:
        lines += ["LANGUAGES", "---------", "  " + ", ".join(langs), ""]
    if d.get("projects"):
        lines += ["PROJECTS", "--------"]
        for p in d["projects"]:
            lines.append(f"  ▸ {p.get('name','')}")
            if p.get("description"): lines.append(f"    {p['description']}")
            if p.get("technologies"): lines.append(f"    Tech: {p['technologies']}")
            lines.append("")
    if d.get("years_of_experience","").strip():
        lines += [f"Years of experience: {d['years_of_experience'].strip()}", ""]
    return "\n".join(lines)

# ── Gradio handler ─────────────────────────────────────────────────────────────
def gradio_extract(file_obj, min_px, max_px, max_new_tokens):
    if file_obj is None:
        return "⚠️  Please upload a file first.", "", ""
    filename = os.path.basename(file_obj).lower()
    if not filename.endswith(ALLOWED_EXTENSIONS):
        return f"⚠️  Unsupported file type: {filename}", "", ""
    try:
        with open(file_obj, "rb") as f:
            raw_bytes = f.read()
        resume = pipeline_from_bytes(raw_bytes, filename, int(min_px), int(max_px), int(max_new_tokens))
        d = resume.model_dump()
        return format_resume(d), json.dumps(d, indent=2, ensure_ascii=False), ""
    except Exception as e:
        return f"❌  {e}", "", str(e)

# ── Gradio UI ──────────────────────────────────────────────────────────────────
with gr.Blocks(title="Resume OCR — Local") as gradio_ui:
    gr.HTML("""
    <div style="text-align:center;padding:1.2rem 0 0.2rem">
      <h1>📄 Resume OCR — Local</h1>
      <p style="color:#6b7280;font-size:0.92rem">
        Extracted locally via Qwen2.5-VL-3B-Instruct &nbsp;·&nbsp;
        REST API at <code>/extract-resume</code> &nbsp;·&nbsp;
        Docs at <code>/docs</code>
      </p>
    </div>
    """)
    with gr.Row(equal_height=False):
        with gr.Column(scale=1, min_width=300):
            file_input = gr.File(
                label="Resume file (JPG / PNG / WEBP / PDF)",
                file_types=[".jpg", ".jpeg", ".png", ".webp", ".pdf"],
            )
            run_btn = gr.Button("Extract →", variant="primary", size="lg")
            with gr.Accordion("⚙️ Advanced settings", open=False):
                max_new_tokens_sl = gr.Slider(512, 4096, value=DEFAULT_MAX_TOK, step=256,
                    label="Max new tokens")
                min_pixels_sl = gr.Slider(64, 512, value=DEFAULT_MIN_PX, step=64,
                    label="Min pixels (×28²)")
                max_pixels_sl = gr.Slider(128, 1280, value=DEFAULT_MAX_PX, step=64,
                    label="Max pixels (×28²)")
        with gr.Column(scale=2, min_width=400):
            with gr.Tabs():
                with gr.Tab("Formatted"):
                    formatted_out = gr.Textbox(label="Parsed resume", lines=30,
                        placeholder="Output will appear here …")
                with gr.Tab("Raw JSON"):
                    json_out = gr.Textbox(label="Raw JSON", lines=30,
                        placeholder="Raw JSON will appear here …")
                with gr.Tab("Debug"):
                    debug_out = gr.Textbox(label="Raw model output (on error only)", lines=20)
    run_btn.click(
        fn=gradio_extract,
        inputs=[file_input, min_pixels_sl, max_pixels_sl, max_new_tokens_sl],
        outputs=[formatted_out, json_out, debug_out],
    )

# ── Mount Gradio into FastAPI — defines the module-level "app" uvicorn finds ──
app = gr.mount_gradio_app(api, gradio_ui, path="/")

if __name__ == "__main__":
    port = int(os.environ.get("HF_PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port)