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