Spaces:
Running
Running
import gradio as gr | |
import torch | |
from transformers import BertTokenizer, BertForSequenceClassification | |
# Load model and tokenizer | |
model_name = "AventIQ-AI/bert-talentmatchai" | |
tokenizer = BertTokenizer.from_pretrained(model_name) | |
model = BertForSequenceClassification.from_pretrained(model_name, torch_dtype=torch.float16) | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
model.to(device) | |
model.eval() | |
# Label mapping | |
label_mapping = {0: "No Fit", 1: "Low Fit", 2: "Potential Fit", 3: "Good Fit"} | |
def preprocess_text(text, max_length=256): | |
"""Truncate input text to avoid exceeding model limits.""" | |
return " ".join(text.split()[:max_length]) | |
def talent_match(resume, job_description): | |
resume = preprocess_text(resume) | |
job_description = preprocess_text(job_description) | |
input_text = f"Resume: {resume} Job Description: {job_description}" | |
inputs = tokenizer([input_text], padding="max_length", truncation=True, return_tensors="pt").to(device) | |
with torch.no_grad(): | |
outputs = model(**inputs) | |
prediction = outputs.logits.argmax(dim=1).item() | |
return label_mapping[prediction] | |
iface = gr.Interface( | |
fn=talent_match, | |
inputs=[ | |
gr.Textbox(label="📄 Resume", placeholder="Paste the candidate's resume here...", lines=5), | |
gr.Textbox(label="📌 Job Description", placeholder="Paste the job description here...", lines=5) | |
], | |
outputs=gr.Textbox(label="✅ Match Result"), | |
title="🤖 AI-Powered Talent Matching System", | |
description="🔍 Enter a candidate's resume and a job description to check if they are a good match using AI.", | |
theme="compact", | |
allow_flagging="never", | |
examples=[ | |
["Experienced Python developer skilled in machine learning and data science.", "Looking for a Python developer with ML experience."], | |
["Project manager with 5 years in Agile methodologies.", "Seeking a Scrum Master with Agile experience."] | |
], | |
) | |
if __name__ == "__main__": | |
iface.launch() | |