| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """Train a character-level GRU for in-browser email autocomplete on Enron.""" |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import re |
| import time |
| from collections import Counter |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnxruntime as ort |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from datasets import load_dataset |
| from huggingface_hub import HfApi |
| from onnxruntime.quantization import QuantType, quantize_dynamic |
|
|
| |
| PAD, EOS, NEWLINE, FIRST_PRINTABLE = 0, 1, 2, 3 |
| VOCAB_SIZE = FIRST_PRINTABLE + (127 - 32) |
| IGNORE = -100 |
|
|
| ENCODE = np.full(256, 255, dtype=np.uint8) |
| ENCODE[10] = NEWLINE |
| ENCODE[32:127] = np.arange(FIRST_PRINTABLE, VOCAB_SIZE, dtype=np.uint8) |
| DECODE = {NEWLINE: "\n", **{FIRST_PRINTABLE + i: chr(32 + i) for i in range(95)}} |
|
|
| LETTER = np.zeros(VOCAB_SIZE, dtype=bool) |
| for char in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ": |
| LETTER[ENCODE[ord(char)]] = True |
|
|
| KEYBOARD_NEIGHBORS = { |
| "q": "wa", "w": "qes", "e": "wrd", "r": "etf", "t": "ryg", "y": "tuh", "u": "yij", |
| "i": "uok", "o": "ipl", "p": "ol", "a": "qsz", "s": "awdx", "d": "serf", "f": "drgv", |
| "g": "fthb", "h": "gyjn", "j": "hukm", "k": "jil", "l": "kop", "z": "asx", "x": "zsc", |
| "c": "xdv", "v": "cfb", "b": "vgn", "n": "bhm", "m": "njk", |
| } |
| NEIGHBOR_TABLE = np.tile(np.arange(VOCAB_SIZE, dtype=np.uint8)[:, None], (1, 4)) |
| for key, neighbors in KEYBOARD_NEIGHBORS.items(): |
| for slot in range(4): |
| NEIGHBOR_TABLE[ENCODE[ord(key)], slot] = ENCODE[ord(neighbors[slot % len(neighbors)])] |
|
|
| CUT_REPLY = re.compile( |
| r"(?im)^\s*(?:-+\s*original message|-+\s*forwarded by|-+\s*inline attachment|from:\s|to:\s|sent:\s" |
| r"|subject:\s|cc:\s|.*\bwrote:\s*$|_{5,}|-{5,}|={5,}|\*{5,}|\d{1,2}/\d{1,2}/\d{2,4}\s+\d{1,2}:\d{2})" |
| ) |
| SENSITIVE = re.compile( |
| r"(?i)https?://|www\.|\S+@\S+|\(?\b\d{3}\)?[-. ]\d{3}[-. ]\d{4}\b|\b\d{3}-\d{2}-\d{4}\b" |
| ) |
| BULK_SENDER = re.compile(r"(?i)announce|newsletter|mailer|outlook\.team|^all[._]|general\.|chairman|enron\.(?:on|net)") |
| SMART_CHARS = str.maketrans({"‘": "'", "’": "'", "“": '"', "”": '"', "–": "-", "—": "-", "\t": " "}) |
|
|
|
|
| def clean_paragraphs(body: str) -> list[str]: |
| text = body.replace("\r\n", "\n").replace("=20", " ").replace("=\n", "").translate(SMART_CHARS) |
| text = "\n".join(line for line in text.split("\n") if not line.lstrip().startswith(">")) |
| cut = CUT_REPLY.search(text) |
| if cut: |
| text = text[: cut.start()] |
| text = re.sub(r"(?m)^-- *\n[\s\S]*", "", text) |
|
|
| paragraphs = [] |
| for block in re.split(r"\n\s*\n", text): |
| lines = [re.sub(r" +", " ", line).strip() for line in block.split("\n") if line.strip()] |
| if not lines: |
| continue |
| |
| joined = lines[0] |
| for previous, line in zip(lines, lines[1:]): |
| joined += (" " if len(previous) >= 50 else "\n") + line |
| letters = sum(char.isalpha() for char in joined) |
| if SENSITIVE.search(joined) or letters < 0.6 * len(joined): |
| continue |
| paragraphs.append(joined) |
| return paragraphs |
|
|
|
|
| def load_emails(limit_frac: float) -> dict[str, list[np.ndarray]]: |
| dataset = load_dataset("corbt/enron-emails", split="train").select_columns(["from", "body"]) |
| if limit_frac < 1: |
| dataset = dataset.shuffle(seed=0).select(range(int(len(dataset) * limit_frac))) |
|
|
| parsed: list[tuple[str, list[str]]] = [] |
| for row in dataset: |
| sender = (row["from"] or "").lower() |
| if "@enron.com" not in sender or BULK_SENDER.search(sender): |
| continue |
| paragraphs = clean_paragraphs(row["body"] or "") |
| if paragraphs: |
| parsed.append((sender, paragraphs)) |
|
|
| counts = Counter(p.lower() for _, paragraphs in parsed for p in paragraphs if len(p) >= 60) |
|
|
| seen: set[bytes] = set() |
| splits: dict[str, list[np.ndarray]] = {"train": [], "val": [], "test": []} |
| for sender, paragraphs in parsed: |
| text = "\n\n".join(p for p in paragraphs if counts[p.lower()] < 5) |
| if not 30 <= len(text) <= 6000: |
| continue |
| digest = hashlib.sha1(re.sub(r"\s+", " ", text.lower()).encode()).digest() |
| if digest in seen: |
| continue |
| seen.add(digest) |
| bucket = int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100 |
| split = "test" if bucket < 5 else "val" if bucket < 10 else "train" |
| ids = ENCODE[np.frombuffer(text.encode("ascii", "ignore"), dtype=np.uint8)] |
| splits[split].append(np.append(ids[ids != 255], EOS).astype(np.uint8)) |
| return splits |
|
|
|
|
| def concat_shuffled(emails: list[np.ndarray], rng: np.random.Generator) -> np.ndarray: |
| return np.concatenate([emails[i] for i in rng.permutation(len(emails))]) |
|
|
|
|
| def add_typos(clean: np.ndarray, word_rate: float, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]: |
| """Corrupt the input stream while the targets keep predicting the clean next character.""" |
| x = clean[:-1].copy() |
| y = clean[1:].astype(np.int16) |
| is_letter = LETTER[x] |
| starts = np.flatnonzero(is_letter[1:] & ~is_letter[:-1]) + 1 |
| ends = np.flatnonzero(is_letter[:-1] & ~is_letter[1:]) + 1 |
| if len(ends) and len(starts) and ends[0] <= starts[0]: |
| ends = ends[1:] |
| count = min(len(starts), len(ends)) |
| starts, ends = starts[:count], ends[:count] |
|
|
| lengths = ends - starts |
| chosen = (lengths >= 3) & (rng.random(count) < word_rate) |
| word_starts, word_lengths = starts[chosen], lengths[chosen] |
| position = word_starts + 1 + (rng.random(len(word_starts)) * (word_lengths - 2)).astype(np.int64) |
| position = position[position + 1 < len(x)] |
| ops = rng.integers(0, 4, size=len(position)) |
| neighbor = NEIGHBOR_TABLE[x[position], rng.integers(0, 4, size=len(position))] |
|
|
| substitute, swap, insert, delete = (position[ops == op] for op in range(4)) |
| x[substitute] = neighbor[ops == 0] |
| first, second = x[swap].copy(), x[swap + 1].copy() |
| x[swap], x[swap + 1], y[swap] = second, first, IGNORE |
|
|
| repeats = np.ones(len(x), dtype=np.int64) |
| repeats[delete] = 0 |
| repeats[insert] = 2 |
| inserted_at = np.cumsum(repeats)[insert] - 1 |
| x_out, y_out = np.repeat(x, repeats), np.repeat(y, repeats) |
| x_out[inserted_at] = neighbor[ops == 2] |
| return x_out, y_out |
|
|
|
|
| def make_batches(x, y, batch_size, seq_len, rng): |
| offset = int(rng.integers(seq_len)) |
| starts = rng.permutation(np.arange(offset, len(x) - seq_len, seq_len)) |
| span = np.arange(seq_len) |
| for index in range(0, len(starts) - batch_size + 1, batch_size): |
| positions = starts[index : index + batch_size][:, None] + span |
| yield torch.from_numpy(x[positions].astype(np.int64)), torch.from_numpy(y[positions].astype(np.int64)) |
|
|
|
|
| class CharGRU(nn.Module): |
| def __init__(self, hidden: int, embedding: int = 64): |
| super().__init__() |
| self.embedding = nn.Embedding(VOCAB_SIZE, embedding, padding_idx=PAD) |
| self.gru = nn.GRU(embedding, hidden, batch_first=True) |
| self.head = nn.Linear(hidden, VOCAB_SIZE) |
|
|
| def forward(self, ids, state=None): |
| output, state = self.gru(self.embedding(ids), state) |
| return self.head(output), state |
|
|
|
|
| def sequence_loss(logits, targets, burn_in): |
| return F.cross_entropy(logits[:, burn_in:].flatten(0, 1), targets[:, burn_in:].flatten(), ignore_index=IGNORE) |
|
|
|
|
| @torch.no_grad() |
| def evaluate_loss(model, emails, args, device, max_chars=5_000_000): |
| rng = np.random.default_rng(1) |
| clean = concat_shuffled(emails, rng)[:max_chars] |
| x, y = clean[:-1], clean[1:].astype(np.int16) |
| total, batches = 0.0, 0 |
| for inputs, targets in make_batches(x, y, args.batch, args.seq, rng): |
| logits, _ = model(inputs.to(device)) |
| total += sequence_loss(logits, targets.to(device), args.burn_in).item() |
| batches += 1 |
| return total / max(batches, 1) |
|
|
|
|
| def decode(ids) -> str: |
| return "".join(DECODE.get(int(i), "") for i in ids) |
|
|
|
|
| def sample_positions(emails, count, context, horizon, rng): |
| contexts, truths = [], [] |
| attempts = 0 |
| while len(contexts) < count and attempts < count * 200: |
| attempts += 1 |
| email = emails[int(rng.integers(len(emails)))] |
| if len(email) < 64: |
| continue |
| cursor = int(rng.integers(16, len(email) - 16)) |
| previous, following = int(email[cursor - 1]), int(email[cursor]) |
| if not (LETTER[previous] or previous == ENCODE[32]) or not LETTER[following]: |
| continue |
| window = np.full(context, EOS, dtype=np.int64) |
| prefix = email[max(0, cursor - context) : cursor] |
| window[-len(prefix):] = prefix |
| truth = email[cursor : cursor + horizon] |
| ending = np.flatnonzero((truth == EOS) | (truth == NEWLINE)) |
| truths.append(decode(truth[: ending[0]] if len(ending) else truth)) |
| contexts.append(window) |
| return np.stack(contexts), truths |
|
|
|
|
| def build_suggestion(ids, probabilities, threshold, max_words=6): |
| text, confidence, best, words = "", 1.0, "", 0 |
| for token, probability in zip(ids, probabilities): |
| token = int(token) |
| if token in (EOS, NEWLINE, PAD): |
| break |
| char = DECODE[token] |
| if char == " ": |
| if confidence >= threshold: |
| best = text |
| words += 1 |
| if words >= max_words: |
| break |
| confidence *= float(probability) |
| if confidence < threshold: |
| break |
| text += char |
| if char in ".,!?": |
| best = text |
| return best |
|
|
|
|
| def score_suggestions(generated, probabilities, truths, thresholds): |
| report = {} |
| for threshold in thresholds: |
| shown, correct, saved = 0, 0, 0 |
| for ids, probs, truth in zip(generated, probabilities, truths): |
| suggestion = build_suggestion(ids, probs, threshold) |
| if len(suggestion) < 4: |
| continue |
| shown += 1 |
| if truth.startswith(suggestion): |
| correct += 1 |
| saved += len(suggestion) |
| report[str(threshold)] = { |
| "coverage": round(shown / len(truths), 4), |
| "precision": round(correct / shown, 4) if shown else None, |
| "avg_chars_when_correct": round(saved / correct, 2) if correct else None, |
| "chars_saved_per_position": round(saved / len(truths), 3), |
| } |
| return report |
|
|
|
|
| @torch.no_grad() |
| def greedy_torch(model, contexts, steps, device, chunk=2048): |
| all_ids, all_probs = [], [] |
| for index in range(0, len(contexts), chunk): |
| ids = torch.from_numpy(contexts[index : index + chunk]).to(device) |
| logits, state = model(ids) |
| step_logits = logits[:, -1] |
| batch_ids, batch_probs = [], [] |
| for _ in range(steps): |
| probabilities = step_logits.softmax(-1) |
| confidence, token = probabilities.max(-1) |
| batch_ids.append(token.cpu().numpy()) |
| batch_probs.append(confidence.cpu().numpy()) |
| logits, state = model(token[:, None], state) |
| step_logits = logits[:, -1] |
| all_ids.append(np.stack(batch_ids, 1)) |
| all_probs.append(np.stack(batch_probs, 1)) |
| return np.concatenate(all_ids), np.concatenate(all_probs) |
|
|
|
|
| def greedy_onnx(session, contexts, steps, hidden): |
| state = np.zeros((len(contexts), hidden), dtype=np.float32) |
| logits = None |
| for column in range(contexts.shape[1]): |
| logits, state = session.run(None, {"input_id": contexts[:, column], "state": state}) |
| ids, confidences = [], [] |
| for _ in range(steps): |
| probabilities = np.exp(logits - logits.max(-1, keepdims=True)) |
| probabilities /= probabilities.sum(-1, keepdims=True) |
| token = probabilities.argmax(-1) |
| ids.append(token) |
| confidences.append(probabilities[np.arange(len(token)), token]) |
| logits, state = session.run(None, {"input_id": token.astype(np.int64), "state": state}) |
| return np.stack(ids, 1), np.stack(confidences, 1) |
|
|
|
|
| class MatMulLinear(nn.Module): |
| """Linear layer that exports as MatMul so ONNX Runtime dynamic quantization applies.""" |
|
|
| def __init__(self, weight: torch.Tensor, bias: torch.Tensor): |
| super().__init__() |
| self.weight = nn.Parameter(weight.t().clone()) |
| self.bias = nn.Parameter(bias.clone()) |
|
|
| def forward(self, inputs): |
| return inputs.matmul(self.weight) + self.bias |
|
|
|
|
| class SingleStep(nn.Module): |
| """One-character GRU update with weights copied from the trained nn.GRU.""" |
|
|
| def __init__(self, model: CharGRU, hidden: int): |
| super().__init__() |
| self.embedding = model.embedding |
| input_weight, hidden_weight = model.gru.weight_ih_l0.detach(), model.gru.weight_hh_l0.detach() |
| input_bias, hidden_bias = model.gru.bias_ih_l0.detach(), model.gru.bias_hh_l0.detach() |
| layers = [] |
| for gate in range(3): |
| row = slice(gate * hidden, (gate + 1) * hidden) |
| layers.append(MatMulLinear(input_weight[row], input_bias[row])) |
| layers.append(MatMulLinear(hidden_weight[row], hidden_bias[row])) |
| self.gates = nn.ModuleList(layers) |
| self.head = MatMulLinear(model.head.weight.detach(), model.head.bias.detach()) |
|
|
| def forward(self, input_id, state): |
| embedded = self.embedding(input_id) |
| reset = torch.sigmoid(self.gates[0](embedded) + self.gates[1](state)) |
| update = torch.sigmoid(self.gates[2](embedded) + self.gates[3](state)) |
| proposal = torch.tanh(self.gates[4](embedded) + reset * self.gates[5](state)) |
| next_state = (1 - update) * proposal + update * state |
| return self.head(next_state), next_state |
|
|
|
|
| def verify_single_step(model, step, hidden): |
| ids = torch.randint(3, VOCAB_SIZE, (4, 32)) |
| logits, _ = model(ids) |
| state = torch.zeros(4, hidden) |
| outputs = [] |
| for column in range(ids.shape[1]): |
| step_logits, state = step(ids[:, column], state) |
| outputs.append(step_logits) |
| difference = (torch.stack(outputs, 1) - logits).abs().max().item() |
| if difference > 1e-3: |
| raise RuntimeError(f"Exported single-step GRU does not match nn.GRU: max difference {difference}") |
| return difference |
|
|
|
|
| def measure_latency(path, hidden, runs=3000): |
| options = ort.SessionOptions() |
| options.intra_op_num_threads = 1 |
| session = ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"]) |
| state = np.zeros((1, hidden), dtype=np.float32) |
| timings = [] |
| for run in range(runs): |
| token = np.array([FIRST_PRINTABLE + run % 90], dtype=np.int64) |
| started = time.perf_counter() |
| _, state = session.run(None, {"input_id": token, "state": state}) |
| timings.append((time.perf_counter() - started) * 1000) |
| return {"p50_ms": round(float(np.percentile(timings, 50)), 3), "p95_ms": round(float(np.percentile(timings, 95)), 3)} |
|
|
|
|
| def export_and_score(model, args, output, contexts, truths, thresholds): |
| model = model.cpu().eval() |
| step = SingleStep(model, args.hidden).eval() |
| difference = verify_single_step(model, step, args.hidden) |
| print(f"single-step GRU matches nn.GRU (max difference {difference:.2e})", flush=True) |
|
|
| float_path = output / "model.onnx" |
| torch.onnx.export( |
| step, |
| (torch.zeros(1, dtype=torch.long), torch.zeros(1, args.hidden)), |
| float_path, |
| input_names=["input_id", "state"], |
| output_names=["logits", "next_state"], |
| dynamic_axes={"input_id": {0: "batch"}, "state": {0: "batch"}, "logits": {0: "batch"}, "next_state": {0: "batch"}}, |
| opset_version=17, |
| dynamo=False, |
| ) |
| quantized_path = output / "model_quint8.onnx" |
| quantize_dynamic(float_path, quantized_path, weight_type=QuantType.QUInt8, op_types_to_quantize=["MatMul"]) |
|
|
| subset = min(1000, len(contexts)) |
| result = {} |
| for name, path in (("float", float_path), ("quint8", quantized_path)): |
| session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) |
| ids, probabilities = greedy_onnx(session, contexts[:subset], args.horizon, args.hidden) |
| result[name] = { |
| "size_mb": round(path.stat().st_size / 1e6, 2), |
| "native_cpu_single_thread_step": measure_latency(path, args.hidden), |
| "suggestions_1000_positions": score_suggestions(ids, probabilities, truths[:subset], thresholds), |
| } |
| return result |
|
|
|
|
| def show_samples(model, device): |
| prompts = [ |
| "Hi John,\n\nThanks for ", "Please let me know if ", "I will send you ", "Let me kn", "Thnaks for the ", |
| "I look forward to ", "Can we schedule a ", "Sorry for the dela", "Plese find attached ", "We need to discuss ", |
| ] |
| samples = {} |
| for prompt in prompts: |
| window = np.full(128, EOS, dtype=np.int64) |
| encoded = ENCODE[np.frombuffer(prompt.encode("ascii"), dtype=np.uint8)] |
| window[-len(encoded):] = encoded |
| ids, probabilities = greedy_torch(model, window[None], 48, device) |
| samples[prompt] = { |
| "continuation": decode(ids[0]), |
| "suggestion_at_0.2": build_suggestion(ids[0], probabilities[0], 0.2), |
| } |
| print(f"{prompt!r} -> {samples[prompt]}", flush=True) |
| return samples |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--hidden", type=int, default=512) |
| parser.add_argument("--epochs", type=int, default=30) |
| parser.add_argument("--batch", type=int, default=256) |
| parser.add_argument("--seq", type=int, default=256) |
| parser.add_argument("--burn-in", type=int, default=32) |
| parser.add_argument("--lr", type=float, default=1.5e-3) |
| parser.add_argument("--typo-rate", type=float, default=0.08) |
| parser.add_argument("--limit-frac", type=float, default=1.0) |
| parser.add_argument("--eval-positions", type=int, default=5000) |
| parser.add_argument("--context", type=int, default=128) |
| parser.add_argument("--horizon", type=int, default=48) |
| parser.add_argument("--repo", default="cfahlgren1/enron-autocomplete") |
| parser.add_argument("--no-push", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
| torch.manual_seed(0) |
| rng = np.random.default_rng(0) |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| print(f"device={device} {torch.cuda.get_device_name(0) if device.type == 'cuda' else ''} args={vars(args)}", flush=True) |
| output = Path("outputs") |
| output.mkdir(parents=True, exist_ok=True) |
|
|
| started = time.time() |
| splits = load_emails(args.limit_frac) |
| data_stats = {name: {"emails": len(e), "chars": int(sum(map(len, e)))} for name, e in splits.items()} |
| print(f"data ready in {time.time() - started:.0f}s: {data_stats}", flush=True) |
| if not splits["train"] or not splits["val"] or not splits["test"]: |
| raise RuntimeError(f"A split is empty: {data_stats}") |
|
|
| model = CharGRU(args.hidden).to(device) |
| parameters = sum(p.numel() for p in model.parameters()) |
| print(f"parameters={parameters:,}", flush=True) |
|
|
| optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-5) |
| steps_per_epoch = max(data_stats["train"]["chars"] // (args.batch * args.seq), 1) |
| total_steps = args.epochs * steps_per_epoch |
| scheduler = torch.optim.lr_scheduler.LambdaLR( |
| optimizer, lambda step: 0.02 + 0.98 * 0.5 * (1 + np.cos(np.pi * min(step / total_steps, 1))) |
| ) |
|
|
| history = [] |
| best_loss = float("inf") |
| best_path = output / "best.pt" |
| for epoch in range(args.epochs): |
| model.train() |
| epoch_started = time.time() |
| x, y = add_typos(concat_shuffled(splits["train"], rng), args.typo_rate, rng) |
| running, batches, chars = 0.0, 0, 0 |
| for inputs, targets in make_batches(x, y, args.batch, args.seq, rng): |
| logits, _ = model(inputs.to(device, non_blocking=True)) |
| loss = sequence_loss(logits, targets.to(device, non_blocking=True), args.burn_in) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| optimizer.step() |
| scheduler.step() |
| running += loss.item() |
| batches += 1 |
| chars += inputs.numel() |
| if batches == 0: |
| raise RuntimeError("No training batches; reduce --batch or --seq for this data size") |
|
|
| model.eval() |
| val_loss = evaluate_loss(model, splits["val"], args, device) |
| seconds = time.time() - epoch_started |
| record = { |
| "epoch": epoch + 1, |
| "train_loss": round(running / batches, 4), |
| "val_loss": round(val_loss, 4), |
| "seconds": round(seconds, 1), |
| "chars_per_sec": int(chars / seconds), |
| } |
| history.append(record) |
| print(json.dumps(record), flush=True) |
| if val_loss < best_loss: |
| best_loss = val_loss |
| torch.save({"model": model.state_dict(), "args": vars(args)}, best_path) |
|
|
| model.load_state_dict(torch.load(best_path, map_location=device)["model"]) |
| model.eval() |
|
|
| thresholds = [0.05, 0.1, 0.2, 0.35, 0.5, 0.7] |
| eval_rng = np.random.default_rng(42) |
| contexts, truths = sample_positions(splits["test"], args.eval_positions, args.context, args.horizon, eval_rng) |
| ids, probabilities = greedy_torch(model, contexts, args.horizon, device) |
| suggestions = score_suggestions(ids, probabilities, truths, thresholds) |
| print(json.dumps({"held_out_sender_suggestions": suggestions}, indent=1), flush=True) |
|
|
| examples = [] |
| for index in range(min(25, len(contexts))): |
| examples.append({ |
| "context": decode(contexts[index])[-60:], |
| "truth": truths[index], |
| "suggestion_at_0.2": build_suggestion(ids[index], probabilities[index], 0.2), |
| }) |
| samples = show_samples(model, device) |
| exported = export_and_score(model, args, output, contexts, truths, thresholds) |
| print(json.dumps(exported, indent=1), flush=True) |
|
|
| report = { |
| "args": vars(args), |
| "device": str(device), |
| "parameters": parameters, |
| "data": data_stats, |
| "history": history, |
| "best_val_loss": round(best_loss, 4), |
| "test_positions": len(contexts), |
| "held_out_sender_suggestions": suggestions, |
| "exports": exported, |
| "prompt_samples": samples, |
| "held_out_examples": examples, |
| "total_seconds": round(time.time() - started, 1), |
| } |
| (output / "report.json").write_text(json.dumps(report, indent=2) + "\n") |
| (output / "vocab.json").write_text(json.dumps({ |
| "pad": PAD, "eos": EOS, "newline": NEWLINE, "first_printable_id": FIRST_PRINTABLE, |
| "first_printable_codepoint": 32, "size": VOCAB_SIZE, "hidden": args.hidden, |
| }, indent=2) + "\n") |
|
|
| if args.no_push: |
| print(f"Skipping upload; artifacts in {output}", flush=True) |
| return |
| api = HfApi(token=os.environ["HF_TOKEN"]) |
| api.create_repo(args.repo, repo_type="model", private=True, exist_ok=True) |
| api.upload_folder(folder_path=str(output), repo_id=args.repo, repo_type="model") |
| print(f"Uploaded to https://huggingface.co/{args.repo}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
