#!/bin/bash
# SFT-only HF Job — fire and forget. ~30 min on a10g-large, ~$0.50.
# Pushes adapter to HF Hub for demo + submission.
set -e
cd /w
export PYTHONPATH="/w:${PYTHONPATH}"

pip install -q 'huggingface_hub>=0.28'

# Pull data files
mkdir -p data
python -c "from huggingface_hub import hf_hub_download; [hf_hub_download(repo_id='Pratham-math/fathom-code', filename=f'data/{f}', local_dir='/w', token='${HF_TOKEN}') for f in ['train.jsonl','eval.jsonl','sft_traces.jsonl']]"

# Same install pattern that worked for smoke
pip install -q openenv-core fastapi 'uvicorn[standard]' pydantic RestrictedPython tiktoken httpx
pip install -q hydra-core omegaconf wandb tyro
pip install -q transformers==4.56.2 accelerate==1.5.2 peft==0.14.0 bitsandbytes==0.45.1
pip install -q datasets==4.7.0
pip install -q --no-deps trl==1.2.0
pip install -q safetensors sentencepiece einops scipy xxhash protobuf pyyaml fsspec aiohttp dill multiprocess pyarrow requests filelock packaging tokenizers regex tqdm

# Optional Unsloth — falls back to plain HF if it fails
pip install -q --no-deps unsloth==2026.4.8 unsloth-zoo || echo "unsloth skipped, using HF transformers"

# SFT on 0.5B (proven to load via smoke)
python <<'PY'
from hydra import initialize, compose
from train.model_load import load_model_and_tokenizer
from train.sft import run_sft
with initialize(config_path="../configs", version_base="1.3"):
    cfg = compose(config_name="config", overrides=["model=qwen_0_5b_smoke","train=sft"])
m, t = load_model_and_tokenizer(cfg)
print("SFT adapter saved at:", run_sft(cfg, m, t))
PY

# Push adapter to HF Hub
HF_USER=$(python -c "from huggingface_hub import HfApi; print(HfApi(token='${HF_TOKEN}').whoami()['name'])")
python -c "
from huggingface_hub import HfApi
api = HfApi(token='${HF_TOKEN}')
repo = '${HF_USER}/fathom-0.5b-sft'
api.create_repo(repo, repo_type='model', exist_ok=True, private=False)
api.upload_folder(folder_path='outputs/sft_adapter', repo_id=repo, repo_type='model')
print(f'Adapter pushed to https://huggingface.co/{repo}')
"