File size: 2,096 Bytes
2e26e31 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | #!/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()
|