Update app.py

app.py

@@ -7,26 +7,10 @@ from collections import Counter
 # Load NER pipeline
 ner_pipeline = pipeline("ner", model="dslim/bert-base-NER", tokenizer="dslim/bert-base-NER", aggregation_strategy="simple")
 
-# Load Job Category Classifier
-text_classifier = pipeline("text-classification", model="serbog/distilbert-jobCategory_410k")
-
-# Mapping from category code to readable label
-CATEGORY_MAP = {
-    "C1": "Engineering",
-    "C2": "Information Technology",
-    "C3": "Sales & Marketing",
-    "C4": "Accounting & Finance",
-    "C5": "Healthcare",
-    "D1": "Education",
-    "D2": "Human Resources",
-    "E1": "Operations & Logistics",
-    "E2": "Legal",
-    "F1": "Customer Support",
-    "Other": "General / Undefined"
-}
+# Load NEW job classifier with human-readable labels
+text_classifier = pipeline("text-classification", model="tkuye/job-description-classifier")
 
 def clean_resume_text(text):
-    """Clean text by removing unwanted characters and formatting."""
     text = re.sub(r'http\S+', ' ', text)
     text = re.sub(r'#\S+', '', text)
     text = re.sub(r'@\S+', ' ', text)
@@ -35,7 +19,6 @@ def clean_resume_text(text):
     return re.sub(r'\s+', ' ', text).strip()
 
 def extract_resume_text(file):
-    """Extract raw text from uploaded PDF."""
     try:
         reader = PyPDF2.PdfReader(file)
         text = ""
@@ -50,7 +33,6 @@ def extract_resume_text(file):
         return None, f"Error reading PDF: {str(e)}"
 
 def classify_resume_ner(entities):
-    """Basic rule-based NER classification using ORG, LOC, MISC."""
     orgs = [e['word'] for e in entities if e['entity_group'] == 'ORG']
     locs = [e['word'] for e in entities if e['entity_group'] == 'LOC']
     jobs = [e['word'] for e in entities if e['entity_group'] == 'MISC']
@@ -66,7 +48,6 @@ def classify_resume_ner(entities):
     }
 
 def process_resumes(files):
-    """Extract entities and show classification based on NER."""
     all_results = {}
     for file in files:
         file_name = file.name.split("/")[-1]
@@ -82,7 +63,7 @@ def process_resumes(files):
         all_results[file_name] = {
             "Persons": list({e["word"] for e in entities if e["entity_group"] == "PER"}),
             "Organizations": list({e["word"] for e in entities if e["entity_group"] == "ORG"}),
-            "Locations": list({e["word"] for e in entities if e["entity_group"] == "LOC"}),
+            "Locations": list({e["word"] for e in entities if e["entity_group"] == "LOC"]),
             "Other Entities": list({e["word"] for e in entities if e["entity_group"] not in ["PER", "ORG", "LOC"]}),
             "Cleaned_Text": cleaned_text,
             "Classification (NER)": classification
@@ -90,7 +71,6 @@ def process_resumes(files):
     return all_results
 
 def classify_resumes_with_model(files):
-    """Use job category model to classify resume into readable job field."""
     predictions = {}
     for file in files:
         file_name = file.name.split("/")[-1]
@@ -99,20 +79,16 @@ def classify_resumes_with_model(files):
             predictions[file_name] = {"error": error}
             continue
         cleaned_text = clean_resume_text(resume_text)
-        result = text_classifier(cleaned_text[:512])  # Truncate for safety
-        raw_label = result[0]['label']
-        readable_label = CATEGORY_MAP.get(raw_label, "Unknown")
-
+        result = text_classifier(cleaned_text[:512])
         predictions[file_name] = {
-            "Predicted Job Category": readable_label,
-            "Raw Label": raw_label,
+            "Predicted Job Category": result[0]['label'],
             "Confidence Score": round(result[0]['score'], 4)
         }
     return predictions
 
 # Gradio Interface
-with gr.Blocks(title="Resume Analyzer") as demo:
-    gr.Markdown("## 📂 Multi-Resume Entity Extractor & Job Classifier\nUpload multiple PDF resumes. This tool extracts entities using NER and predicts the job field using a trained classifier model.")
+with gr.Blocks(title="Resume Analyzer (Readable Labels)") as demo:
+    gr.Markdown("## 📂 Multi-Resume Entity Extractor & Job Classifier\nUpload multiple PDF resumes. This tool extracts entities using NER and predicts job category using a model with readable labels.")
 
     with gr.Row():
         file_input = gr.File(file_types=[".pdf"], label="Upload Resume PDFs", file_count="multiple")
@@ -122,7 +98,7 @@ with gr.Blocks(title="Resume Analyzer") as demo:
         classify_button = gr.Button("🧠 Predict Job Category (Model)")
 
     output_entities = gr.JSON(label="NER Results & Classification")
-    output_class = gr.JSON(label="Model-Predicted Job Category")
+    output_class = gr.JSON(label="Predicted Job Category (Model)")
 
     extract_button.click(fn=process_resumes, inputs=[file_input], outputs=[output_entities])
     classify_button.click(fn=classify_resumes_with_model, inputs=[file_input], outputs=[output_class])