updated

backend/services/resume_parser.py

@@ -1,104 +1,49 @@
-from __future__ import annotations
-import os
-import re
-import subprocess
-import zipfile
-from typing import List, Dict
-import torch
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
+import subprocess, zipfile, re, os
 
-# ===============================
-# Load Model & Tokenizer
-# ===============================
-MODEL_ID = "sravya-abburi/ResumeParserBERT"
+# === Load pretrained HF model instead of training ===
+MODEL_NAME = "sravya-abburi/ResumeParserBERT"  # or Kiet/autotrain-resume_parser-1159242747
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME)
+ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
 
-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
-)
-
-# ===============================
-# Text Extraction
-# ===============================
+# === Extract text from PDF/DOCX ===
 def extract_text(file_path: str) -> str:
-    """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"):
-            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 ""
-
-# ===============================
-# Parse Resume using BERT NER
-# ===============================
-def parse_with_bert(text: str) -> Dict[str, str]:
-    """Parse resume text into structured fields using BERT NER."""
+    if file_path.lower().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 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 with NER ===
+def parse_resume(file_path: str) -> dict:
+    text = extract_text(file_path)
     entities = ner_pipeline(text)
 
-    name_tokens, skill_tokens, edu_tokens, exp_tokens = [], [], [], []
-
+    name, skills, education, experience = [], [], [], []
     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)
+        word = ent["word"]
+        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(name_tokens),
-        "skills": ", ".join(skill_tokens),
-        "education": ", ".join(edu_tokens),
-        "experience": ", ".join(exp_tokens)
+        "name": " ".join(set(name)),
+        "skills": ", ".join(set(skills)),
+        "education": ", ".join(set(education)),
+        "experience": ", ".join(set(experience))
     }
-
-# ===============================
-# Main Parse Function
-# ===============================
-def parse_resume(file_path: str, filename: str) -> dict:
-    """Main function for resume parsing."""
-    text = extract_text(file_path)
-    if not text:
-        return {"name": "", "skills": "", "education": "", "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