inference.py

"""
Inference script for the fine-tuned Career Agent.
Usage:
    python inference.py --model Builder-Neekhil/career-agent-v1
"""
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="Builder-Neekhil/career-agent-v1", help="HF model repo")
    parser.add_argument("--device", default="auto", help="Device map")
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        torch_dtype=torch.bfloat16,
        device_map=args.device,
        trust_remote_code=True,
    )

    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)

    messages = [
        {"role": "system", "content": (
            "You are a seasoned career advising expert with 15 years of experience. "
            "Be specific, honest, actionable, and concise."
        )},
        {"role": "user", "content": "Review my resume for a software engineering role."},
    ]

    out = pipe(messages, return_full_text=False)
    print(out[0]["generated_text"])


if __name__ == "__main__":
    main()