import argparse
from src.data_loader import load_data
from src.model import load_model
from src.trainer import train_model
from src.evaluate import evaluate_model
from src.inference import infer_resume
from transformers import AutoTokenizer

def run():
    parser = argparse.ArgumentParser(description="ATS Resume Optimizer")
    parser.add_argument("mode", choices=["train", "evaluate", "infer"], help="Mode to run the script")
    parser.add_argument("--resume", type=str, help="Path to resume for inference")
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
    datasets = load_data(tokenizer)
    model = load_model()

    if args.mode == "train":
        train_model(model, datasets)
    elif args.mode == "evaluate":
        evaluate_model(model, datasets)
    elif args.mode == "infer":
        if not args.resume:
            raise ValueError("Resume path must be provided in infer mode.")
        infer_resume(model, tokenizer, args.resume)

if __name__ == "__main__":
    run()