| |
| import argparse |
| import re |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| from mla.checkpoint import load_checkpoint |
| from mla.tokenizer import Tokenizer |
| from mla.loss import cross_entropy |
| from mla.chat import ChatSession |
|
|
| SFT = Path("data/sft") |
| TOK = Path("data/tokenizer/tokenizer.json") |
| _CODE = re.compile(r"`|</|/>|def |import |class |printf|console\.log|#include|<tool|print\(|for \(|;\s*$") |
|
|
|
|
| def get_sft_batch(ids, mask, block, batch, rng): |
| hi = len(ids) - block - 1 |
| ix = rng.integers(0, hi, size=batch) |
| x = np.stack([ids[i:i + block] for i in ix]).astype(np.int64) |
| y = np.stack([ids[i + 1:i + 1 + block] for i in ix]).astype(np.int64) |
| ym = np.stack([mask[i + 1:i + 1 + block] for i in ix]) |
| y[ym == 0] = -1 |
| return x, y |
|
|
|
|
| def eval_loss_acc(model, ids, mask, block, batch, n_batches, rng): |
| tot_loss = 0.0 |
| correct = 0 |
| total = 0 |
| for _ in range(n_batches): |
| x, y = get_sft_batch(ids, mask, block, batch, rng) |
| out = model(x) |
| tot_loss += float(cross_entropy(out, y).data) |
| pred = np.asarray(out.data).argmax(axis=-1) |
| keep = (y != -1) |
| correct += int((pred[keep] == y[keep]).sum()) |
| total += int(keep.sum()) |
| mean = tot_loss / max(1, n_batches) |
| return mean, float(np.exp(mean)), correct / max(1, total) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", type=str, default="checkpoints/sft_final.npz") |
| ap.add_argument("--block-size", type=int, default=256) |
| ap.add_argument("--batch-size", type=int, default=32) |
| ap.add_argument("--eval-batches", type=int, default=40) |
| ap.add_argument("--seed", type=int, default=123) |
| args = ap.parse_args() |
|
|
| tok = Tokenizer.load(TOK) |
| model, _, step = load_checkpoint(args.ckpt) |
| val_ids = np.load(SFT / "val_ids.npy") |
| val_mask = np.load(SFT / "val_mask.npy") |
| rng = np.random.default_rng(args.seed) |
|
|
| loss, ppl, acc = eval_loss_acc(model, val_ids, val_mask, args.block_size, |
| args.batch_size, args.eval_batches, rng) |
| print(f"ckpt={args.ckpt} step={step}") |
| print(f"[assistant-masked] val_loss={loss:.4f} ppl={ppl:.2f} next_token_acc={acc:.3f}") |
|
|
| print("\n[sample chat turns]") |
| chat_prompts = [ |
| "I had a really rough day at work today.", |
| "I'm feeling lonely lately.", |
| "my dog passed away last week.", |
| "I just got a promotion, I'm so excited!", |
| "hey, how are you?", |
| ] |
| s = ChatSession(model, tok, temperature=0.8, top_k=40, top_p=0.9, seed=7) |
| for u in chat_prompts: |
| print(f" USER: {u}") |
| print(f" BOT : {s.reply(u)}") |
|
|
| print("\n[refusal / in-scope check]") |
| code_prompts = [ |
| "Write a Python function to reverse a string.", |
| "Give me the code for bubble sort in C++.", |
| "How do I import numpy and print hello world?", |
| ] |
| passed = 0 |
| for u in code_prompts: |
| cs = ChatSession(model, tok, temperature=0.7, top_k=40, top_p=0.9, seed=3) |
| r = cs.reply(u) |
| emitted_code = bool(_CODE.search(r)) |
| ok = not emitted_code |
| passed += ok |
| print(f" USER: {u}") |
| print(f" BOT : {r}") |
| print(f" in_scope={'PASS' if ok else 'FAIL (emitted code-like text)'}") |
| print(f"\nrefusal: {passed}/{len(code_prompts)} stayed in-scope") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
