File size: 5,437 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | #!/usr/bin/env python3
"""A/B картинками: родной энкодер MOSS (Qwen3.5-9B) против 0.8B + адаптер.
Один промпт, один сид, один размер — картинки встают рядом, слева родной, справа адаптер.
Условие родного энкодера должно быть посчитано заранее (`native_embeds.py`): 15.6 ГБ учителя и
DiT в 16 ГБ не влезают одновременно, а условие от промпта не зависит.
python3 src/native_embeds.py --out out/native_embeds.pt
python3 src/sample_ab.py --adapter out/adapter_v1.safetensors --out-dir out/samples
"""
import argparse
import os
import sys
import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw, ImageFont
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from adapter_lib import load_adapter # noqa: E402
from moss_lib import (MOSS_DIR, Student, build_pipeline, condition, # noqa: E402
empty_cache, generate, load_tokenizer)
from prompts import AB # noqa: E402
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"
def grid(rows, out_path, gap=8, header=40, caption=28, bg=(24, 24, 24), fg=(235, 235, 235)):
"""rows: [(подпись, prompt, [PIL, PIL])]; колонки = родной / адаптер."""
try:
font_h, font_c = ImageFont.truetype(FONT, 20), ImageFont.truetype(FONT, 14)
except OSError:
font_h = font_c = ImageFont.load_default()
width = max(sum(im.width for im in r[2]) + gap * (len(r[2]) - 1) for r in rows)
height = header + sum(caption + gap + max(im.height for im in r[2]) for r in rows)
canvas = Image.new("RGB", (width, height), bg)
draw = ImageDraw.Draw(canvas)
draw.text((8, 10), "родной Qwen3.5-9B (слева) / 0.8B + адаптер (справа)", font=font_h, fill=fg)
y = header
for name, prompt, images in rows:
draw.text((8, y), f"{name}: {prompt[:110]}", font=font_c, fill=fg)
y += caption
x = 0
for im in images:
canvas.paste(im, (x, y))
x += im.width + gap
y += max(im.height for im in images) + gap
canvas.save(out_path)
return out_path
def cosine(a, b):
n = min(a.shape[0], b.shape[0])
return F.cosine_similarity(a[:n].float(), b[:n].float(), dim=-1).mean().item()
def main():
ap = argparse.ArgumentParser(description="A/B: родной энкодер MOSS против адаптера")
ap.add_argument("--adapter", default=os.path.join(ROOT, "adapter.safetensors"))
ap.add_argument("--native", default=os.path.join(ROOT, "out", "native_embeds.pt"))
ap.add_argument("--moss-dir", default=MOSS_DIR)
ap.add_argument("--student", default=None,
help="по умолчанию локальная копия в <moss-dir>/text_encoder")
ap.add_argument("--out-dir", default=os.path.join(ROOT, "out", "samples"))
ap.add_argument("--size", type=int, default=1024)
ap.add_argument("--steps", type=int, default=30)
ap.add_argument("--guidance", type=float, default=3.0)
ap.add_argument("--seed", type=int, default=1234)
args = ap.parse_args()
if not os.path.exists(args.native):
raise SystemExit(f"нет {args.native} — сначала запустите src/native_embeds.py")
native = torch.load(args.native, map_location="cpu", weights_only=True)
os.makedirs(args.out_dir, exist_ok=True)
tok = load_tokenizer(args.moss_dir)
adapter, meta = load_adapter(args.adapter, device="cuda")
student = Student(args.student, [int(x) for x in str(meta["student_layers"]).split(",")])
pos = dict(zip([n for n, _ in AB], condition(adapter, student, tok, [t for _, t in AB])))
neg = condition(adapter, student, tok, [""])[0]
for name, emb in pos.items():
print(f" {name:10s} cos(условие адаптера, родное) = {cosine(emb[0], native['items'][name]):.3f}",
flush=True)
del student, adapter
empty_cache()
pipe = build_pipeline(args.moss_dir)
pipe.to("cuda")
rows = []
for name, prompt in AB:
adapter_img = generate(pipe, pos[name], neg, size=args.size, steps=args.steps,
seed=args.seed, guidance=args.guidance)
nat_emb = (native["items"][name], torch.ones(native["items"][name].shape[0], dtype=torch.bool))
nat_img = generate(pipe, nat_emb, (native["items"]["negative"],
torch.ones(native["items"]["negative"].shape[0], dtype=torch.bool)),
size=args.size, steps=args.steps, seed=args.seed, guidance=args.guidance)
adapter_img.save(os.path.join(args.out_dir, f"{name}_adapter.png"))
nat_img.save(os.path.join(args.out_dir, f"{name}_native.png"))
rows.append((name, prompt, [nat_img, adapter_img]))
print(f" {name}: готово", flush=True)
out = grid(rows, os.path.join(args.out_dir, "ab_native_vs_adapter.png"))
print(f"готово: {out}, пик VRAM {torch.cuda.max_memory_allocated() / 2**30:.1f} ГиБ", flush=True)
if __name__ == "__main__":
main()
|