Update app.py

app.py

@@ -4,11 +4,13 @@ import gradio as gr
 from transformers import pipeline
 from collections import Counter
 
-# Load the Hugging Face NER pipeline
+# Load NER pipeline
 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="MoritzLaurer/bert-multilingual-passage-reranking-msmarco")
+
 def clean_resume_text(text):
-    """Clean resume 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)
@@ -17,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 file."""
     try:
         reader = PyPDF2.PdfReader(file)
         text = ""
@@ -31,8 +32,7 @@ def extract_resume_text(file):
     except Exception as e:
         return None, f"Error reading PDF: {str(e)}"
 
-def classify_resume(entities):
-    """Classify resume based on dominant entity types."""
+def classify_resume_ner(entities):
     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']
@@ -44,45 +44,65 @@ def classify_resume(entities):
     return {
         "Main_Organization": dominant_org[0][0] if dominant_org else "Unknown",
         "Main_Location": dominant_loc[0][0] if dominant_loc else "Unknown",
-        "Possible_Job/Field": dominant_job[0][0] if dominant_job else "General"
+        "Possible_Job/Field (NER)": dominant_job[0][0] if dominant_job else "General"
     }
 
-def extract_entities_from_pdfs(files):
-    """Process multiple resumes, extract entities, and classify."""
-    summary = {}
-
+def process_resumes(files):
+    all_results = {}
     for file in files:
         file_name = file.name.split("/")[-1]
         resume_text, error = extract_resume_text(file)
 
         if error:
-            summary[file_name] = {"error": error}
+            all_results[file_name] = {"error": error}
             continue
 
         cleaned_text = clean_resume_text(resume_text)
         entities = ner_pipeline(cleaned_text)
+        classification = classify_resume_ner(entities)
 
-        result = {
+        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"}),
             "Other": list({e["word"] for e in entities if e["entity_group"] not in ["PER", "ORG", "LOC"]}),
             "Cleaned_Text": cleaned_text,
-            "Classification": classify_resume(entities)
+            "Classification (NER)": classification
         }
+    return all_results
 
-        summary[file_name] = result
-
-    return summary
+def classify_resumes_with_model(files):
+    predictions = {}
+    for file in files:
+        file_name = file.name.split("/")[-1]
+        resume_text, error = extract_resume_text(file)
+        if error:
+            predictions[file_name] = {"error": error}
+            continue
+        cleaned_text = clean_resume_text(resume_text)
+        result = text_classifier(cleaned_text[:512])  # Truncate long resumes
+        predictions[file_name] = {
+            "Predicted Label (HuggingFace Classifier)": result[0]['label'],
+            "Confidence": round(result[0]['score'], 4)
+        }
+    return predictions
 
 # Gradio UI
-iface = gr.Interface(
-    fn=extract_entities_from_pdfs,
-    inputs=gr.File(file_types=[".pdf"], label="Upload Resumes (PDF)", file_count="multiple"),
-    outputs=gr.JSON(label="Resume Classification & Entity Summary"),
-    title="📂 Multi-Resume Entity Extractor & Classifier",
-    description="Upload multiple PDF resumes. This tool extracts text, identifies key entities, and classifies each resume by organizations, locations, and possible job/field."
-)
+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.")
+
+    with gr.Row():
+        file_input = gr.File(file_types=[".pdf"], label="Upload Resume PDFs", file_count="multiple")
+
+    with gr.Row():
+        extract_button = gr.Button("🔍 Extract & Analyze Entities")
+        classify_button = gr.Button("🧠 Predict Job Role with Classifier")
+
+    output_entities = gr.JSON(label="Entity Extraction & NER Classification")
+    output_class = gr.JSON(label="Predicted Job Classification (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])
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()