ayushsinha's picture
Upload 2 files
7621713 verified
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()