"""Upload the 135M checkpoints and their demo material to the private assets repo.

Uploads weight-only checkpoints (the optimizer state is stripped, halving the EP file), the
tokenizer, the sample notebook, the training-curve figure, and a card describing how the pair was
produced. The repo stays private; making a public model release is a separate decision.
"""
from pathlib import Path

from huggingface_hub import HfApi

REPO = 'blackhao0426/ept-assets'
RUN = Path('/home/yurenh2/ept/ep_run')
ASSETS = Path('/home/yurenh2/ept/assets')
PREFIX = 'models/fw135m'

CARD = """# 135M transformer language models: Equilibrium Propagation and its backprop twin

Two checkpoints of the same architecture, trained from scratch on FineWeb-Edu, differing only in the
training rule.

| file | training rule | val CE (tail mean, last 10%) | best val CE |
|---|---|---|---|
| `ep_fw135m_s440000.pt` | Equilibrium Propagation, no backward pass anywhere | 3.2087 | 3.0423 |
| `bp_fw135m_s440000.pt` | backpropagation (matched twin) | 3.2071 | 3.0902 |

A second backprop seed reached 3.2076, so the two backprop seeds differ by 0.0005 and EP sits 0.0013
above their mean, which is 0.1% in perplexity. With one EP seed and two backprop seeds this design
cannot resolve a difference of that size; more seeds are running.

Architecture: 12 layers, width 768, 12 heads, context 256, OLMo2-style blocks, untied 32k-vocabulary
output. 440k steps at effective batch 24, Muon hybrid optimizer, cosine schedule, mixed precision.
Trained on 2.7B tokens.

Contents are weight-only: `tok`, `pos`, `blocks`, `wout`, `lnf`, plus `step`, `val`, and the full
training `config`. Optimizer state is stripped. Both files load with the same code, since the two
training scripts save identical keys.

Sampling is an ordinary forward pass; Equilibrium Propagation appears only during training. See
`EPT_135M_samples.ipynb` for generations from both models on identical prompts and seeds, and
`fig_135m_curves.png` for the validation curves, including a third run that plateaus 45% higher in
perplexity because its contrast readout recovered the nudge by subtracting two large states, which
destroys the part of it that falls below single-precision resolution.
"""


def main():
    api = HfApi()
    items = [
        (RUN / 'runs/share_fw135m_ep.pt', f'{PREFIX}/ep_fw135m_s440000.pt'),
        (RUN / 'runs/share_fw135m_bp.pt', f'{PREFIX}/bp_fw135m_s440000.pt'),
        (RUN / 'data/fineweb_edu/tokenizer.json', f'{PREFIX}/tokenizer.json'),
        (ASSETS / 'EPT_135M_samples.ipynb', f'{PREFIX}/EPT_135M_samples.ipynb'),
        (ASSETS / 'figs/fig_135m_curves.png', f'{PREFIX}/fig_135m_curves.png'),
    ]
    card = RUN / 'runs/_card_135m.md'
    card.write_text(CARD)
    items.append((card, f'{PREFIX}/README.md'))

    for src, dst in items:
        if not src.exists():
            print(f'MISSING {src}')
            continue
        api.upload_file(path_or_fileobj=str(src), path_in_repo=dst, repo_id=REPO,
                        repo_type='dataset')
        print(f'uploaded {dst}  ({src.stat().st_size/1e6:.1f} MB)')
    print(f'\nhttps://huggingface.co/datasets/{REPO}/tree/main/{PREFIX}')


if __name__ == '__main__':
    main()