arittrabag's picture
Update app.py
02ca209 verified
import gradio as gr
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
# Load fine-tuned model & tokenizer
model_path = "bert_resume_classifier" # Change if saved elsewhere
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
# Label mapping
label_map = {
0: "Advocate", 1: "Arts", 2: "Automation Testing", 3: "Blockchain",
4: "Business Analyst", 5: "Civil Engineer", 6: "Data Science", 7: "Database",
8: "DevOps Engineer", 9: "DotNet Developer", 10: "ETL Developer",
11: "Electrical Engineering", 12: "HR", 13: "Hadoop", 14: "Health and fitness",
15: "Java Developer", 16: "Mechanical Engineer", 17: "Network Security Engineer",
18: "Operations Manager", 19: "PMO", 20: "Python Developer",
21: "SAP Developer", 22: "Sales", 23: "Testing", 24: "Web Designing"
}
# Prediction Function
def predict_resume_category(resume_text):
inputs = tokenizer(resume_text, truncation=True, padding=True, max_length=512, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predicted_class = torch.argmax(logits, dim=1).item()
return f"Predicted Job Category: {label_map[predicted_class]}"
# Gradio Interface
iface = gr.Interface(
fn=predict_resume_category,
inputs=gr.Textbox(lines=10, placeholder="Paste resume text here..."),
outputs="text",
title="Resume Job Category Predictor",
description="Enter resume text to classify the job category using BERT.",
)
# Launch
iface.launch(share=True) # Use share=True to get a public Gradio link