Spaces:
Paused
Paused
from __future__ import annotations | |
import os | |
import re | |
import subprocess | |
import zipfile | |
import json | |
import torch | |
from typing import List | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
# =============================== | |
# Load DeepSeek Janus-Pro-7B Model | |
# =============================== | |
MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.2" | |
print(f"Loading {MODEL_ID}... (This may take some time on first run)") | |
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) | |
model = AutoModelForCausalLM.from_pretrained( | |
MODEL_ID, | |
torch_dtype=torch.float16, | |
device_map="auto" | |
) | |
# =============================== | |
# Text Extraction (PDF/DOCX) | |
# =============================== | |
def extract_text(file_path: str) -> str: | |
"""Extract text from PDF or DOCX resumes.""" | |
if not file_path or not os.path.isfile(file_path): | |
return "" | |
lower_name = file_path.lower() | |
try: | |
if lower_name.endswith('.pdf'): | |
result = subprocess.run( | |
['pdftotext', '-layout', file_path, '-'], | |
stdout=subprocess.PIPE, | |
stderr=subprocess.PIPE, | |
check=False | |
) | |
return result.stdout.decode('utf-8', errors='ignore') | |
elif lower_name.endswith('.docx'): | |
with zipfile.ZipFile(file_path) as zf: | |
with zf.open('word/document.xml') as docx_xml: | |
xml_bytes = docx_xml.read() | |
xml_text = xml_bytes.decode('utf-8', errors='ignore') | |
xml_text = re.sub(r'<w:p[^>]*>', '\n', xml_text, flags=re.I) | |
text = re.sub(r'<[^>]+>', ' ', xml_text) | |
return re.sub(r'\s+', ' ', text) | |
else: | |
return "" | |
except Exception: | |
return "" | |
# =============================== | |
# Name Extraction (Fallback) | |
# =============================== | |
def extract_name(text: str, filename: str) -> str: | |
"""Extract candidate's name from resume text or filename.""" | |
if text: | |
lines = [ln.strip() for ln in text.splitlines() if ln.strip()] | |
for line in lines[:10]: | |
if re.match(r'(?i)resume|curriculum vitae', line): | |
continue | |
words = line.split() | |
if 1 < len(words) <= 4: | |
if all(re.match(r'^[A-ZÀ-ÖØ-Þ][\w\-]*', w) for w in words): | |
return line | |
base = os.path.basename(filename) | |
base = re.sub(r'\.(pdf|docx|doc)$', '', base, flags=re.I) | |
base = re.sub(r'[\._-]+', ' ', base) | |
base = re.sub(r'(?i)\b(cv|resume)\b', '', base) | |
return base.title().strip() | |
# =============================== | |
# Janus-Pro Parsing | |
# =============================== | |
def parse_with_deepseek(text: str) -> dict: | |
"""Use DeepSeek Janus-Pro-7B to extract resume details in JSON format.""" | |
prompt = f""" | |
Extract the following information from the resume text below: | |
- Full Name | |
- Skills (comma separated) | |
- Education (degrees + institutions) | |
- Experience (job titles + companies) | |
Return only valid JSON in the following structure: | |
{{ | |
"name": "Full Name", | |
"skills": "Skill1, Skill2, Skill3", | |
"education": "Degree1 - Institution1; Degree2 - Institution2", | |
"experience": "Job1 - Company1; Job2 - Company2" | |
}} | |
Resume: | |
{text} | |
""" | |
inputs = tokenizer(prompt, return_tensors="pt").to(model.device) | |
outputs = model.generate(**inputs, max_new_tokens=512) | |
response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
# Extract JSON safely | |
match = re.search(r"\{.*\}", response, re.S) | |
if match: | |
try: | |
return json.loads(match.group()) | |
except: | |
pass | |
return {"name": "", "skills": "", "education": "", "experience": ""} | |
# =============================== | |
# Fallback Heading-based Parsing | |
# =============================== | |
def fallback_parse(text: str) -> dict: | |
"""Simple heading-based parsing as backup.""" | |
skills = re.findall(r"Skills\s*[:\-]?\s*(.*)", text, re.I) | |
education = re.findall(r"Education\s*[:\-]?\s*(.*)", text, re.I) | |
experience = re.findall(r"(Experience|Work History)\s*[:\-]?\s*(.*)", text, re.I) | |
return { | |
"skills": ", ".join(skills), | |
"education": ", ".join(education), | |
"experience": ", ".join([exp[1] for exp in experience]) | |
} | |
# =============================== | |
# Main Parse Function | |
# =============================== | |
def parse_resume(file_path: str, filename: str) -> dict: | |
"""Main resume parsing function.""" | |
text = extract_text(file_path) | |
name = extract_name(text, filename) | |
# Try Janus-Pro parsing | |
ents = parse_with_deepseek(text) | |
# If Janus-Pro misses fields, use fallback | |
if not ents.get("skills") or not ents.get("education"): | |
fb = fallback_parse(text) | |
ents["skills"] = ents.get("skills") or fb["skills"] | |
ents["education"] = ents.get("education") or fb["education"] | |
ents["experience"] = ents.get("experience") or fb["experience"] | |
return { | |
"name": ents.get("name") or name, | |
"skills": ents.get("skills", ""), | |
"education": ents.get("education", ""), | |
"experience": ents.get("experience", "") | |
} | |