""" 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("""

📄 Resume OCR — Local

Extracted locally via Qwen2.5-VL-3B-Instruct  ·  REST API at /extract-resume  ·  Docs at /docs

""") 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)