updated

backend/services/resume_parser.py

@@ -3,100 +3,102 @@ import os
 import re
 import subprocess
 import zipfile
-from typing import List
-from transformers import pipeline
+from typing import List, Dict
+import torch
+from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 
 # ===============================
-# Load Lightweight Resume Parser
+# Load Model & Tokenizer
 # ===============================
-resume_parser_model = pipeline("text-classification", model="Kiet/autotrain-resume_parser-1159242747")
+MODEL_ID = "sravya-abburi/ResumeParserBERT"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForTokenClassification.from_pretrained(MODEL_ID)
+
+ner_pipeline = pipeline(
+    "ner",
+    model=model,
+    tokenizer=tokenizer,
+    aggregation_strategy="simple",
+    device=0 if torch.cuda.is_available() else -1
+)
 
 # ===============================
-# PDF/DOCX Text Extraction
+# Text Extraction
 # ===============================
 def extract_text(file_path: str) -> str:
-    """Extract text from PDF or DOCX resumes."""
+    """Extract raw text from PDF or DOCX."""
     if not file_path or not os.path.isfile(file_path):
         return ""
 
     lower_name = file_path.lower()
     try:
-        if lower_name.endswith('.pdf'):
+        if lower_name.endswith(".pdf"):
             result = subprocess.run(
-                ['pdftotext', '-layout', file_path, '-'],
+                ["pdftotext", "-layout", file_path, "-"],
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
                 check=False
             )
-            return result.stdout.decode('utf-8', errors='ignore')
+            return result.stdout.decode("utf-8", errors="ignore")
 
-        elif lower_name.endswith('.docx'):
+        elif lower_name.endswith(".docx"):
             with zipfile.ZipFile(file_path) as zf:
-                with zf.open('word/document.xml') as docx_xml:
+                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)
+                    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 ""
 
 # ===============================
-# Fallback Name Extraction
+# Parse Resume using BERT NER
 # ===============================
-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()
+def parse_with_bert(text: str) -> Dict[str, str]:
+    """Parse resume text into structured fields using BERT NER."""
+    entities = ner_pipeline(text)
 
-# ===============================
-# Model-based Resume Parsing
-# ===============================
-def parse_with_kiet_model(text: str) -> dict:
-    """Use Kiet's resume parser model to extract fields."""
-    try:
-        # The pipeline might return structured text (needs post-processing)
-        parsed_output = resume_parser_model(text)
-        
-        # Since the model output may vary, we simulate structured mapping
-        return {
-            "name": parsed_output[0]['label'] if parsed_output else "",
-            "skills": "Extracted Skills Here",
-            "education": "Extracted Education Here",
-            "experience": "Extracted Experience Here"
-        }
-    except Exception:
-        return {"name": "", "skills": "", "education": "", "experience": ""}
+    name_tokens, skill_tokens, edu_tokens, exp_tokens = [], [], [], []
+
+    for ent in entities:
+        label = ent["entity_group"].upper()
+        word = ent["word"].strip()
+
+        if label == "NAME" and word not in name_tokens:
+            name_tokens.append(word)
+        elif label == "SKILL" and word not in skill_tokens:
+            skill_tokens.append(word)
+        elif label == "EDUCATION" and word not in edu_tokens:
+            edu_tokens.append(word)
+        elif label == "EXPERIENCE" and word not in exp_tokens:
+            exp_tokens.append(word)
+
+    return {
+        "name": " ".join(name_tokens),
+        "skills": ", ".join(skill_tokens),
+        "education": ", ".join(edu_tokens),
+        "experience": ", ".join(exp_tokens)
+    }
 
 # ===============================
 # Main Parse Function
 # ===============================
 def parse_resume(file_path: str, filename: str) -> dict:
-    """Main function to parse resumes."""
+    """Main function for resume parsing."""
     text = extract_text(file_path)
-    name = extract_name(text, filename)
-    
-    ents = parse_with_kiet_model(text)
-    if not ents.get("name"):
-        ents["name"] = name
+    if not text:
+        return {"name": "", "skills": "", "education": "", "experience": ""}
 
-    return {
-        "name": ents.get("name", ""),
-        "skills": ents.get("skills", ""),
-        "education": ents.get("education", ""),
-        "experience": ents.get("experience", "")
-    }
+    ents = parse_with_bert(text)
+
+    # Fallback: use filename for name if model doesn't find one
+    if not ents["name"]:
+        base = os.path.basename(filename)
+        base = re.sub(r"\.(pdf|docx|doc)$", "", base, flags=re.I)
+        ents["name"] = re.sub(r"[\._-]+", " ", base).title().strip()
+
+    return ents