import os
import sys
import argparse


def main():
    parser = argparse.ArgumentParser(description="Download GPT-2 via kagglehub to target directory")
    parser.add_argument("--target", type=str, default="/home/azureuser/code/REPST/GPT-2", help="Target directory to store GPT-2")
    args = parser.parse_args()

    try:
        import kagglehub
    except Exception as e:
        print("[ERROR] kagglehub 未安装或导入失败。请先运行: pip install kagglehub")
        print("       需在 ~/.kaggle/kaggle.json 配置 Kaggle API。")
        sys.exit(1)

    os.makedirs(args.target, exist_ok=True)

    handle = "openai/gpt-2"
    print(f"开始通过 kagglehub 下载 {handle} 到 {args.target} ...")
    try:
        path = kagglehub.model_download(handle, path=args.target)
    except Exception as e:
        print(f"[ERROR] 下载失败: {e}")
        sys.exit(2)

    print(f"下载完成，已保存到: {path}")


if __name__ == "__main__":
    main()