| |
| """Ориентир для 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 |
| from prompts import AB |
|
|
| 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() |
|
|
