Spaces:
Paused
Paused
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline | |
import zipfile, re, os | |
from PyPDF2 import PdfReader # Lightweight & already in Spaces | |
# =============================== | |
# Load Model & Tokenizer | |
# =============================== | |
MODEL_NAME = "sravya-abburi/ResumeParserBERT" # Swap to Kiet model if needed | |
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME) | |
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple") | |
# =============================== | |
# Extract Text (PDF & DOCX) | |
# =============================== | |
def extract_text(file_path: str) -> str: | |
"""Extract text from PDF or DOCX without external dependencies.""" | |
file_path_lower = file_path.lower() | |
# β PDF reading using PyPDF2 (no fitz, no installs needed) | |
if file_path_lower.endswith(".pdf"): | |
text = "" | |
with open(file_path, "rb") as f: | |
reader = PdfReader(f) | |
for page in reader.pages: | |
page_text = page.extract_text() | |
if page_text: | |
text += page_text + "\n" | |
return text | |
# β DOCX reading by extracting XML content | |
elif file_path_lower.endswith(".docx"): | |
with zipfile.ZipFile(file_path) as zf: | |
with zf.open("word/document.xml") as docx_xml: | |
xml_text = docx_xml.read().decode("utf-8", errors="ignore") | |
xml_text = re.sub(r"<w:p[^>]*>", "\n", xml_text, flags=re.I) | |
return re.sub(r"<[^>]+>", " ", xml_text) | |
return "" | |
# =============================== | |
# Parse Resume | |
# =============================== | |
def parse_resume(file_path: str, filename: str = None) -> dict: | |
"""Parse resume and extract structured information.""" | |
text = extract_text(file_path) | |
entities = ner_pipeline(text) | |
name, skills, education, experience = [], [], [], [] | |
for ent in entities: | |
label = ent["entity_group"].upper() | |
word = ent["word"].strip() | |
if label == "NAME": | |
name.append(word) | |
elif label == "SKILL": | |
skills.append(word) | |
elif label in ["EDUCATION", "DEGREE"]: | |
education.append(word) | |
elif label in ["EXPERIENCE", "JOB", "ROLE"]: | |
experience.append(word) | |
return { | |
"name": " ".join(dict.fromkeys(name)), | |
"skills": ", ".join(dict.fromkeys(skills)), | |
"education": ", ".join(dict.fromkeys(education)), | |
"experience": ", ".join(dict.fromkeys(experience)) | |
} | |