Spaces:
Paused
Paused
| import os | |
| import re | |
| import subprocess | |
| import zipfile | |
| import json | |
| from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline | |
| # -------------------- | |
| # Load Model | |
| # -------------------- | |
| MODEL_NAME = "sravya-abburi/ResumeParserBERT" | |
| 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 | |
| # -------------------- | |
| def extract_text(file_path: str) -> str: | |
| """Extract text from PDF/DOCX resumes.""" | |
| if file_path.lower().endswith(".pdf"): | |
| try: | |
| result = subprocess.run( | |
| ["pdftotext", "-layout", file_path, "-"], | |
| stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False | |
| ) | |
| return result.stdout.decode("utf-8", errors="ignore") | |
| except: | |
| return "" | |
| elif file_path.lower().endswith(".docx"): | |
| try: | |
| 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) | |
| except: | |
| return "" | |
| return "" | |
| # -------------------- | |
| # Parse Resume | |
| # -------------------- | |
| def parse_resume(file_path: str, filename: str = None) -> dict: | |
| """Extract Name, Skills, Education, Experience from resume.""" | |
| 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)) | |
| } | |
| # -------------------- | |
| # Example | |
| # -------------------- | |
| if __name__ == "__main__": | |
| resume_path = "resume.pdf" # Change to test file | |
| result = parse_resume(resume_path) | |
| print(json.dumps(result, indent=2)) | |