goldphish2209/multilabel-skill-classifier

1import gradio as gr2import onnxruntime as rt3from transformers import AutoTokenizer4import torch5import json6import numpy as np7 8 9# Load tokenizer10tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")11 12# Load skill mapping13with open("skill_mapping.json", "r") as f:14    skill_id = json.load(f)15 16skills = list(skill_id.keys())17 18# Load ONNX model19inf_session = rt.InferenceSession('skill-classifier.onnx')20 21def classify_job_skills(job_description, threshold=0.5):22    """23    Classify skills from a job description24    25    Args:26        job_description: Text of the job posting27        threshold: Minimum confidence score (0-1)28    29    Returns:30        Dictionary of skill -> probability for skills above threshold31    """32    if not job_description.strip():33        return {}34    35    # Tokenize input with attention_mask36    inputs = tokenizer(37        job_description, 38        truncation=True, 39        max_length=512,40        padding='max_length',41        return_tensors='np'42    )43    44    # Run inference with both input_ids and attention_mask45    logits = inf_session.run(46        None,  # Get all outputs47        {48            'input_ids': inputs['input_ids'].astype(np.int64),49            'attention_mask': inputs['attention_mask'].astype(np.int64)50        }51    )[0]52    53    # Convert to probabilities54    probs = torch.sigmoid(torch.FloatTensor(logits))[0]55    56    # Filter by threshold and return top skills57    results = {58        skill: float(prob) 59        for skill, prob in zip(skills, probs) 60        if prob >= threshold61    }62    63    # Sort by probability (highest first)64    return dict(sorted(results.items(), key=lambda x: x[1], reverse=True))65 66# Example job descriptions67examples = [68    [69        """We're looking for a Senior Machine Learning Engineer to join our team. 70        Responsibilities include building ML pipelines, training deep learning models, 71        and deploying models to production using AWS and Docker. Strong Python skills required, 72        along with experience in PyTorch or TensorFlow. Knowledge of MLOps practices and 73        CI/CD pipelines is a plus.""",74        0.575    ],76    [77        """Full Stack Developer needed! Must have strong JavaScript, React, and Node.js experience. 78        You'll be building responsive web applications, working with REST APIs, and managing 79        databases (SQL/NoSQL). Familiarity with Git, Docker, and cloud platforms (AWS/Azure) 80        is required. Great communication and teamwork skills essential.""",81        0.582    ],83    [84        """Data Analyst position available. Looking for someone skilled in SQL, Python, and 85        Excel for data analysis and visualization. Experience with Tableau or Power BI required. 86        You'll perform EDA, create dashboards, and communicate insights to stakeholders. 87        Strong attention to detail and problem-solving skills needed.""",88        0.589    ]90]91 92# Create Gradio interface93with gr.Blocks(title="Job Skills Classifier") as iface:94    gr.Markdown(95        """96        # 🎯 Job Skills Classifier97        98        Extract required skills from job descriptions using AI. 99        Paste a job posting below and click "Classify Skills" to see the detected skills100        """101    )102    103    with gr.Row():104        with gr.Column():105            job_input = gr.Textbox(106                lines=8,107                placeholder="Paste a job description here...",108                label="Job Description"109            )110            threshold_slider = gr.Slider(111                minimum=0.1,112                maximum=0.9,113                value=0.5,114                step=0.05,115                label="Confidence Threshold",116                info="Only show skills with probability above this value"117            )118            classify_btn = gr.Button("Classify Skills", variant="primary")119        120        with gr.Column():121            output_label = gr.Label(122                num_top_classes=20,123                label="Detected Skills"124            )125    126    gr.Markdown("### 💡 Try these examples:")127    gr.Examples(128        examples=examples,129        inputs=[job_input, threshold_slider],130        outputs=output_label,131        fn=classify_job_skills,132        cache_examples=False133    )134    135    gr.Markdown(136        """137        ---138        ### 📊 About139        140        This model detects **technical skills** (Python, Machine Learning, AWS, etc.) and 141        **soft skills** (Communication, Leadership, Problem Solving, etc.) from job descriptions.142        143        **Skills covered:** 80+ technical and soft skills across software development, 144        data science, cloud computing, and more from the tech field145        146        """147    )148    149    # Connect button to function150    classify_btn.click(151        fn=classify_job_skills,152        inputs=[job_input, threshold_slider],153        outputs=output_label154    )155 156# Launch157iface.launch()