app.py

import os
import re
import tempfile
import uuid
from pathlib import Path

import gradio as gr
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    pipeline,
    AutoModelForSeq2SeqLM,
)
from sentence_transformers import SentenceTransformer, util
from pdfminer.high_level import extract_text
import nltk

# ------------------------------------------------------------------
# 1. One-time NLTK resource
# ------------------------------------------------------------------
try:
    nltk.data.find("tokenizers/punkt")
except LookupError:
    nltk.download("punkt")

# ------------------------------------------------------------------
# 2. Global model / pipeline setup (cached on first run)
# ------------------------------------------------------------------
device = 0 if torch.cuda.is_available() else -1

# NER
ner_pipe = pipeline(
    "ner",
    model="dslim/bert-base-NER",
    tokenizer="dslim/bert-base-NER",
    aggregation_strategy="simple",
    device=device,
)

# Sentence similarity
sbert = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# LLM for suggestions
sugg_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
sugg_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
if device >= 0:
    sugg_model = sugg_model.to("cuda")

# Multilingual translation (NLLB-200)
trans_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
trans_model = AutoModelForSeq2SeqLM.from_pretrained(
    "facebook/nllb-200-distilled-600M"
)

# ------------------------------------------------------------------
# 3. Helper functions
# ------------------------------------------------------------------
def translate_to_english(text: str) -> str:
    """
    Detects language via heuristics and translates to English using NLLB-200.
    """
    try:
        text.encode("ascii")
        return text
    except UnicodeEncodeError:
        pass

    inputs = trans_tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=512,
    )
    if device >= 0:
        inputs = {k: v.to("cuda") for k, v in inputs.items()}

    bos_token_id = trans_tokenizer.convert_tokens_to_ids("eng_Latn")
    translated = trans_model.generate(
        **inputs,
        forced_bos_token_id=bos_token_id,
        max_length=512,
    )
    return trans_tokenizer.batch_decode(translated, skip_special_tokens=True)[0]


def extract_text_from_file(file_path: str) -> str:
    path = Path(file_path)
    if path.suffix.lower() == ".pdf":
        return extract_text(file_path)
    with open(file_path, encoding="utf-8") as f:
        return f.read()


def extract_skills_and_experience(resume_text: str):
    ner_results = ner_pipe(resume_text)
    entities = {ent["word"].strip() for ent in ner_results}
    skill_pattern = re.compile(
        r"\b(Python|SQL|Java|C\+\+|JavaScript|React|AWS|Docker|Kubernetes|TensorFlow|PyTorch|Scikit-learn|Pandas|Numpy|Excel|PowerBI|Tableau|SEO|Photoshop|Figma|CAD|SolidWorks|MATLAB|R|Go|Node\.js|Vue|Angular|MongoDB|MySQL|PostgreSQL|GCP|Azure|Linux|Git|CI/CD|Agile|Scrum|Machine Learning|Deep Learning|NLP|Computer Vision|Data Engineering|Data Science|Marketing|Sales|Design|UI|UX|DevOps|SRE|Frontend|Backend|Full-stack)\b",
        flags=re.I,
    )
    regex_skills = set(skill_pattern.findall(resume_text))
    skills = entities.union(regex_skills)
    return list(skills)


def classify_domain(skills: list[str]) -> str:
    domain_keywords = {
        "Data Science": {"python", "sql", "pandas", "numpy", "machine learning"},
        "Software Engineering": {"java", "c++", "python", "javascript", "react"},
        "Marketing": {"seo", "google ads", "facebook ads", "content marketing"},
        "Design": {"photoshop", "figma", "ui", "ux"},
        "DevOps": {"docker", "kubernetes", "aws", "ci/cd"},
    }
    skill_set = {s.lower() for s in skills}
    best_domain = "General"
    best_score = 0
    for domain, keywords in domain_keywords.items():
        score = len(skill_set & keywords)
        if score > best_score:
            best_score = score
            best_domain = domain
    return best_domain


def suggest_job_titles(domain: str, skills: list[str], top_k: int = 5) -> list[str]:
    job_bank = {
        "Data Science": [
            "Data Scientist",
            "Machine Learning Engineer",
            "Data Analyst",
            "AI Researcher",
            "Business Intelligence Analyst",
        ],
        "Software Engineering": [
            "Backend Developer",
            "Frontend Developer",
            "Full-stack Developer",
            "Software Engineer",
            "DevOps Engineer",
        ],
        "Marketing": [
            "Digital Marketing Specialist",
            "SEO Specialist",
            "Content Marketing Manager",
            "Growth Hacker",
            "Social Media Manager",
        ],
        "Design": [
            "UI/UX Designer",
            "Product Designer",
            "Graphic Designer",
            "Web Designer",
            "Motion Designer",
        ],
        "DevOps": [
            "DevOps Engineer",
            "Site Reliability Engineer",
            "Cloud Engineer",
            "Platform Engineer",
            "System Administrator",
        ],
        "General": [
            "Project Manager",
            "Product Manager",
            "Consultant",
            "Technical Writer",
            "Operations Analyst",
        ],
    }
    pool = job_bank.get(domain, job_bank["General"])
    skill_text = " ".join(skills).lower()
    skill_emb = sbert.encode(skill_text, convert_to_tensor=True)
    job_embs = sbert.encode(pool, convert_to_tensor=True)
    scores = util.cos_sim(skill_emb, job_embs)[0]
    top_indices = torch.topk(scores, k=min(top_k, len(pool))).indices
    return [pool[i] for i in top_indices]


def generate_recommendations(skills: list[str], domain: str, job: str) -> str:
    prompt = (
        f"Give concise advice for a candidate applying for {job} in {domain}. "
        f"They currently have these skills: {', '.join(skills)}. "
        f"What skills are missing and what should they improve?"
    )
    inputs = sugg_tokenizer.encode(prompt, return_tensors="pt")
    if device >= 0:
        inputs = inputs.to("cuda")
    outputs = sugg_model.generate(
        inputs, max_length=150, num_beams=3, early_stopping=True
    )
    return sugg_tokenizer.decode(outputs[0], skip_special_tokens=True)


def build_report(
    skills: list[str],
    domain: str,
    job_titles: list[str],
    recommendations: str,
):
    report_lines = ["=== ResumeRanker AI Report ===", ""]
    report_lines.append("Extracted Skills:")
    report_lines.extend(f"- {s}" for s in skills)
    report_lines.append("")
    report_lines.append(f"Inferred Career Domain: {domain}")
    report_lines.append("")
    report_lines.append("Suggested Job Titles:")
    report_lines.extend(f"- {j}" for j in job_titles)
    report_lines.append("")
    report_lines.append("Recommendations / Skill Gaps:")
    report_lines.append(recommendations)
    return "\n".join(report_lines)


# ------------------------------------------------------------------
# 4. Gradio UI
# ------------------------------------------------------------------
def process_resume(file_obj):
    with tempfile.NamedTemporaryFile(
        delete=False, suffix=Path(file_obj.name).suffix
    ) as tmp:
        tmp.write(file_obj.read())
        tmp_path = tmp.name

    raw_text = extract_text_from_file(tmp_path)
    english_text = translate_to_english(raw_text)

    skills = extract_skills_and_experience(english_text)
    domain = classify_domain(skills)
    job_titles = suggest_job_titles(domain, skills)
    recommendations = generate_recommendations(skills, domain, job_titles[0])
    report = build_report(skills, domain, job_titles, recommendations)

    os.remove(tmp_path)

    return (
        ", ".join(skills),
        domain,
        ", ".join(job_titles),
        recommendations,
        report,
    )


# ------------------------------------------------------------------
# 5. Launch
# ------------------------------------------------------------------
if __name__ == "__main__":
    demo = gr.Interface(
        fn=process_resume,
        inputs=gr.File(
            label="Upload Resume (PDF or TXT)",
            file_types=[".pdf", ".txt"],
            type="binary",
        ),
        outputs=[
            gr.Textbox(label="Extracted Skills"),
            gr.Textbox(label="Career Domain"),
            gr.Textbox(label="Suggested Job Titles"),
            gr.Textbox(label="Recommendations"),
            gr.Textbox(label="Full Report"),
        ],
        title="ResumeRanker AI",
        description="Upload your resume and get smart AI-powered job matching & improvement suggestions.",
        allow_flagging="never",
    )

    demo.launch()