Update app.py

app.py

@@ -4,14 +4,14 @@ import gradio as gr
 from transformers import pipeline
 from collections import Counter
 
-# Load NER pipeline
+# Load NER pipeline for entity extraction
 ner_pipeline = pipeline("ner", model="dslim/bert-base-NER", tokenizer="dslim/bert-base-NER", aggregation_strategy="simple")
 
-# Load text classification model (replace with a job-role classifier if available)
-text_classifier = pipeline("text-classification", model="khaimait/job-title-classification-bert")
-
+# Load Job Category Classifier
+text_classifier = pipeline("text-classification", model="serbog/distilbert-jobCategory_410k")
 
 def clean_resume_text(text):
+    """Clean text by removing URLs, punctuation, non-ASCII chars."""
     text = re.sub(r'http\S+', ' ', text)
     text = re.sub(r'#\S+', '', text)
     text = re.sub(r'@\S+', ' ', text)
@@ -20,6 +20,7 @@ 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 = ""
@@ -34,6 +35,7 @@ def extract_resume_text(file):
         return None, f"Error reading PDF: {str(e)}"
 
 def classify_resume_ner(entities):
+    """Classify by extracting key orgs and locations from NER output."""
     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']
@@ -49,6 +51,7 @@ def classify_resume_ner(entities):
     }
 
 def process_resumes(files):
+    """Process multiple resumes with NER and classification."""
     all_results = {}
     for file in files:
         file_name = file.name.split("/")[-1]
@@ -66,13 +69,14 @@ def process_resumes(files):
             "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"}),
-            "Other": list({e["word"] for e in entities if e["entity_group"] not in ["PER", "ORG", "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
         }
     return all_results
 
 def classify_resumes_with_model(files):
+    """Use job category model to predict the field/role."""
     predictions = {}
     for file in files:
         file_name = file.name.split("/")[-1]
@@ -81,26 +85,26 @@ 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 long resumes
+        result = text_classifier(cleaned_text[:512])  # Truncate if too long
         predictions[file_name] = {
-            "Predicted Label (HuggingFace Classifier)": result[0]['label'],
-            "Confidence": round(result[0]['score'], 4)
+            "Predicted Job Category": result[0]['label'].replace("_", " ").title(),
+            "Confidence Score": round(result[0]['score'], 4)
         }
     return predictions
 
-# Gradio UI
-with gr.Blocks(title="Multi-Resume Entity & Job Classifier") as demo:
-    gr.Markdown("## 📂 Multi-Resume Entity Extractor & Classifier\nUpload multiple PDF resumes below. This tool extracts text, identifies key entities, and classifies job field using a Hugging Face model.")
+# Gradio Interface
+with gr.Blocks(title="Resume Analyzer") as demo:
+    gr.Markdown("## 📂 Multi-Resume Entity Extractor & Job Category Classifier\nUpload multiple PDF resumes. This tool uses NER to extract info and a job classification model to predict job field/category.")
 
     with gr.Row():
-        file_input = gr.File(file_types=[".pdf"], label="Upload Resume PDFs", file_count="multiple")
+        file_input = gr.File(file_types=[".pdf"], label="Upload Resumes (PDF)", file_count="multiple")
 
     with gr.Row():
-        extract_button = gr.Button("🔍 Extract & Analyze Entities")
-        classify_button = gr.Button("🧠 Predict Job Role with Classifier")
+        extract_button = gr.Button("🔍 Extract Entities")
+        classify_button = gr.Button("🧠 Predict Job Category")
 
-    output_entities = gr.JSON(label="Entity Extraction & NER Classification")
-    output_class = gr.JSON(label="Predicted Job Classification (Model)")
+    output_entities = gr.JSON(label="Entity Results & NER Classification")
+    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])