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()