utils/generator.py

"""
generator.py — AI Resume Improvement Generator.
 
Uses google/flan-t5-base (free, open-source) to generate:
  1. An improved version of the resume text
  2. A polished professional summary
 
FLAN-T5 is instruction-tuned, making it ideal for text rewriting tasks
without requiring fine-tuning.
 
Model size: ~250 MB. Downloaded once and cached automatically by
HuggingFace transformers.
"""
 
# ---------------------------------------------------------------------------
# Lazy model loading — keeps Streamlit startup fast
# ---------------------------------------------------------------------------
_pipeline = None
 
 
def _get_pipeline():
    """Load and cache the FLAN-T5 text2text pipeline."""
    global _pipeline
    if _pipeline is None:
        from transformers import pipeline
        _pipeline = pipeline(
            "text2text-generation",
            model="google/flan-t5-base",
            max_new_tokens=512,
        )
    return _pipeline
 
 
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
 
def generate_improved_resume(
    resume_text: str,
    job_description: str = "",
    missing_sections: list = None,
) -> str:
    """
    Generate an improved, ATS-friendly version of the resume.
 
    Rewrites the resume with:
      - Professional tone
      - Bullet-point format
      - Stronger action verbs
      - Explicit mention of any missing sections
 
    Args:
        resume_text     : original extracted resume text
        job_description : optional job description to tailor the rewrite
        missing_sections: list of sections the model should add
 
    Returns:
        Improved resume text as a string.
    """
    if not resume_text:
        return "No resume text provided for improvement."
 
    if missing_sections is None:
        missing_sections = []
 
    pipe = _get_pipeline()
 
    # Build a focused, instruction-style prompt
    # FLAN-T5 works best with clear task instructions
    truncated_resume = _truncate(resume_text, 400)
    jd_hint = f"\nThe target job requires: {_truncate(job_description, 100)}" if job_description else ""
    missing_hint = (
        f"\nMake sure to include sections for: {', '.join(missing_sections)}."
        if missing_sections
        else ""
    )
 
    prompt = (
        f"Rewrite this resume to be professional, ATS-friendly, and impactful. "
        f"Use bullet points, strong action verbs, and quantify achievements. "
        f"Keep all original facts.{jd_hint}{missing_hint}\n\n"
        f"Original Resume:\n{truncated_resume}\n\n"
        f"Improved Resume:"
    )
 
    try:
        result = pipe(prompt, max_new_tokens=512, do_sample=False)
        improved = result[0]["generated_text"].strip()
        return improved if improved else "⚠️ Could not generate improvement. Please try again."
    except Exception as e:
        return f"⚠️ Generation error: {str(e)}"
 
 
def generate_professional_summary(
    name: str,
    skills: list,
    experience_present: bool,
    job_title: str = "",
) -> str:
    """
    Generate a short professional summary / objective paragraph.
 
    Args:
        name              : candidate name (or empty string)
        skills            : list of extracted skill strings
        experience_present: whether work experience was detected
        job_title         : optional target job title
 
    Returns:
        A 2–3 sentence professional summary.
    """
    pipe = _get_pipeline()
 
    skills_str = ", ".join(skills[:8]) if skills else "various technical skills"
    exp_phrase = (
        "with hands-on professional experience"
        if experience_present
        else "seeking to start a professional career"
    )
    target = f" targeting a {job_title} role" if job_title else ""
 
    prompt = (
        f"Write a 2-3 sentence professional resume summary for a candidate "
        f"{exp_phrase} in {skills_str}{target}. "
        f"Keep it concise, first-person, and ATS-friendly."
    )
 
    try:
        result = pipe(prompt, max_new_tokens=120, do_sample=False)
        return result[0]["generated_text"].strip()
    except Exception as e:
        return f"⚠️ Summary generation error: {str(e)}"
 
 
def generate_section_content(section_name: str, context: str) -> str:
    """
    Generate content for a specific missing resume section.
 
    Args:
        section_name: e.g., 'Projects', 'Skills', 'Experience'
        context     : relevant resume text to base the generation on
 
    Returns:
        Generated section content as a string.
    """
    pipe = _get_pipeline()
 
    prompt = (
        f"Based on the following resume context, write a professional "
        f"'{section_name}' section in bullet-point format:\n\n"
        f"{_truncate(context, 200)}\n\n"
        f"{section_name} Section:"
    )
 
    try:
        result = pipe(prompt, max_new_tokens=200, do_sample=False)
        return result[0]["generated_text"].strip()
    except Exception as e:
        return f"⚠️ Section generation error: {str(e)}"
 
 
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
 
def _truncate(text: str, max_words: int) -> str:
    """Truncate text to max_words words to stay within model context limit."""
    words = text.split()
    if len(words) > max_words:
        return " ".join(words[:max_words]) + "..."
    return text