File size: 8,545 Bytes
2c87024 | 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """Read-only FP32 prefix continuation, without training or a KV cache."""
import argparse
import json
import time
from pathlib import Path
import sentencepiece as spm
import torch
from vimeml.training.data import file_sha, write_json
from vimeml.training.model import GPTConfig, TinyGPT
ROOT = Path(__file__).resolve().parents[3]
def filtered_logits(logits, forbidden, temperature=1.0, top_k=0, top_p=1.0):
if temperature <= 0 or top_k < 0 or not 0 < top_p <= 1:
raise ValueError("Invalid sampling parameters.")
logits = logits.float().clone() / temperature
logits[list(forbidden)] = -torch.inf
if top_k:
cutoff = torch.topk(logits, min(top_k, logits.numel())).values[-1]
logits[logits < cutoff] = -torch.inf
if top_p < 1:
sorted_values, order = torch.sort(logits, descending=True)
probabilities = torch.softmax(sorted_values, dim=-1)
# Keep the token that crosses the threshold, including the first token.
removed = probabilities.cumsum(-1) - probabilities >= top_p
logits[order[removed]] = -torch.inf
return logits
class JapaneseLM:
def __init__(self, checkpoint, tokenizer_dir, device="cpu"):
self.device = torch.device(device)
if self.device.type not in {"cpu", "cuda"}:
raise ValueError("Choose CPU or CUDA for this first inference tool.")
if self.device.type == "cuda" and not torch.cuda.is_available():
raise ValueError("CUDA unavailable.")
checkpoint, tokenizer_dir = Path(checkpoint), Path(tokenizer_dir)
saved = torch.load(checkpoint, map_location="cpu", weights_only=True)
if saved.get("format") != "vimeml_tiny_gpt_v1":
raise ValueError("Unsupported checkpoint format.")
token_manifest_path = ROOT / saved["config"]["token_dir"] / "manifest.json"
if file_sha(token_manifest_path) != saved["signatures"]["tokens"]:
raise ValueError("Token manifest does not match checkpoint.")
token_manifest = json.loads(token_manifest_path.read_text(encoding="utf-8"))
expected = [digest for path, digest in token_manifest["input_sha256"].items()
if path.replace("\\", "/").rsplit("/", 1)[-1] == "tokenizer.model"]
model_path = tokenizer_dir / "tokenizer.model"
if len(expected) != 1 or file_sha(model_path) != expected[0]:
raise ValueError("Tokenizer model does not match checkpoint vocabulary.")
self.processor = spm.SentencePieceProcessor(model_file=str(model_path))
self.special = {name: getattr(self.processor, f"{name}_id")() for name in ("pad", "unk", "bos", "eos")}
if self.special != token_manifest["special_ids"]:
raise ValueError("Special token IDs do not match training.")
self.model = TinyGPT(GPTConfig(**saved["model_config"]))
if self.processor.vocab_size() != self.model.config.vocab_size:
raise ValueError("Vocabulary size mismatch.")
self.model.load_state_dict(saved["model"])
self.model.to(self.device).eval()
self.metadata = {"checkpoint": str(checkpoint.resolve()), "checkpoint_sha256": file_sha(checkpoint),
"tokenizer_sha256": expected[0], "checkpoint_step": saved["step"],
"model": self.model.configuration(), "device": str(self.device), "precision": "fp32",
"policy": "Plain prefix continuation, separate sentence; no sliding context or KV cache."}
self.forbidden = tuple(self.special[name] for name in ("pad", "unk", "bos"))
del saved
def prefix_ids(self, prompt):
content = self.processor.encode(prompt, out_type=int)
if self.processor.decode(content) != prompt:
raise ValueError("Prompt does not roundtrip with this tokenizer.")
ids = [self.special["bos"], *content]
if len(ids) > self.model.config.context_length:
raise ValueError("Prompt exceeds context length; shorten it rather than silently truncating.")
return ids
@torch.inference_mode()
def next_logits(self, ids):
inputs = torch.tensor([ids], device=self.device, dtype=torch.long)
return self.model(inputs)[0, -1].float()
def next_pieces(self, prompt, count=10):
logits = filtered_logits(self.next_logits(self.prefix_ids(prompt)), self.forbidden)
probabilities = logits.softmax(-1)
values, indices = torch.topk(probabilities, min(count, probabilities.numel() - len(self.forbidden)))
return [{"id": int(index), "piece": self.processor.id_to_piece(int(index)), "probability": float(value)}
for index, value in zip(indices, values)]
def generate(self, prompt, max_new_tokens=32, temperature=0.8, top_k=50, top_p=0.9, seed=42):
if max_new_tokens < 1 or temperature < 0 or top_k < 0 or not 0 < top_p <= 1:
raise ValueError("Invalid generation parameters.")
prefix = self.prefix_ids(prompt)
ids = list(prefix)
generator = torch.Generator(device=self.device).manual_seed(seed)
stop = "max_new_tokens"
for _ in range(max_new_tokens):
if len(ids) > self.model.config.context_length:
stop = "context_limit"
break
logits = self.next_logits(ids)
if temperature == 0:
logits = filtered_logits(logits, self.forbidden)
token = int(logits.argmax())
else:
logits = filtered_logits(logits, self.forbidden, temperature, top_k, top_p)
token = int(torch.multinomial(logits.softmax(-1), 1, generator=generator))
ids.append(token)
if token == self.special["eos"]:
stop = "eos"
break
text = self.processor.decode(ids[1:])
return {"prompt": prompt, "text": text, "new_token_ids": ids[len(prefix):],
"stop_reason": stop, "temperature": temperature, "top_k": top_k, "top_p": top_p,
"seed": seed, "contains_replacement_character": "\ufffd" in text,
"prefix_preserved": text.startswith(prompt)}
def main(argv=None):
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", type=Path, default=ROOT / "artifacts/models/tiny-ja-v1/best.pt")
parser.add_argument("--tokenizer", type=Path, default=ROOT / "artifacts/tokenizers/ja-unigram-16k-v1")
parser.add_argument("--prompts", type=Path, default=ROOT / "configs/inference-prompts.json")
parser.add_argument("--prompt", help="Single custom prefix, instead of the fixed suite.")
parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
parser.add_argument("--threads", type=int, default=4)
parser.add_argument("--max-new-tokens", type=int, default=32)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--output", type=Path, default=ROOT / "outputs/model-checks/tiny-ja-v1")
args = parser.parse_args(argv)
if args.threads < 1:
parser.error("threads must be positive.")
torch.set_num_threads(args.threads)
lm = JapaneseLM(args.checkpoint, args.tokenizer, args.device)
cases = [{"kind": "custom", "prompt": args.prompt}] if args.prompt is not None else json.loads(args.prompts.read_text(encoding="utf-8"))
results, paragraphs = [], []
started = time.perf_counter()
for index, case in enumerate(cases):
prompt = case["prompt"]
greedy = lm.generate(prompt, args.max_new_tokens, temperature=0, seed=args.seed)
sampled = lm.generate(prompt, args.max_new_tokens, seed=args.seed + index)
results.append({**case, "next_pieces": lm.next_pieces(prompt), "greedy": greedy, "sampled": sampled})
paragraph = f"Prompt: {prompt}\nGreedy: {greedy['text']}\nSampled: {sampled['text']}\n"
paragraphs.append(paragraph)
print(paragraph, flush=True)
args.output.mkdir(parents=True, exist_ok=True)
write_json(args.output / "samples.json", {"metadata": lm.metadata, "max_new_tokens": args.max_new_tokens,
"elapsed_seconds": time.perf_counter() - started,
"note": "Qualitative developer prompts, not held-out metrics or an iOS latency benchmark. Top pieces are subword tokens, not phrase candidates.",
"cases": results})
(args.output / "samples.txt").write_text("\n".join(paragraphs), encoding="utf-8")
print(f"Samples: {args.output / 'samples.json'}")
if __name__ == "__main__":
main()
|