#!/usr/bin/env python3
"""Ориентир для A/B: условие от РОДНОГО энкодера (Qwen3.5-9B) на фиксированный набор промптов.

Учитель (15.63 ГиБ) и DiT (9.5 ГиБ) в 16 ГБ не живут вместе, поэтому условие считается один раз
и кладётся на диск; `sample_ab.py` потом берёт его оттуда.

  python3 src/native_embeds.py --out out/native_embeds.pt
"""
import argparse
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from moss_lib import MOSS_DIR, encode, load_teacher, load_tokenizer   # noqa: E402
from prompts import AB                                                # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def main():
    ap = argparse.ArgumentParser(description="условие родного энкодера MOSS для A/B")
    ap.add_argument("--moss-dir", default=MOSS_DIR)
    ap.add_argument("--out", default=os.path.join(ROOT, "out", "native_embeds.pt"))
    ap.add_argument("--teacher-dir", default=None)
    ap.add_argument("--prefer-8bit", action="store_true")
    args = ap.parse_args()

    tok = load_tokenizer(args.moss_dir)
    texts = [t for _, t in AB] + [""]
    rows = encode(tok, texts, 4096)
    teacher, mode = load_teacher(prefer_8bit=args.prefer_8bit, path=args.teacher_dir)

    items = {}
    for (name, _), ids in zip(AB + [("negative", "")], rows):
        y = teacher.last_hidden(torch.tensor([ids], device="cuda"),
                                torch.ones(1, len(ids), dtype=torch.long, device="cuda"))
        items[name] = y[0].to(torch.bfloat16).cpu()
        print(f"  {name:10s} {len(ids):3d} токенов, норма {y[0].norm(dim=-1).mean():.0f}", flush=True)

    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    torch.save({"items": items, "mode": mode, "prompts": dict(AB)}, args.out)
    print(f"готово: {args.out} (учитель {mode})", flush=True)


if __name__ == "__main__":
    main()