CLEF dual-task proper-protocol FT (seed 20261101, epoch 3) from jev-retrain MI355X run 2026-11-04
073de96 verified | """ClefHost: one request through the host steps around the graphs (the graphs themselves are the caller's callables). | |
| host = ClefHost.from_dir("<folder>") # tokenizer.json, embed_tokens.safetensors, lm_head.safetensors, | |
| # joint_head.safetensors, joint_head_config.json | |
| enc = host.encode(request, tile=448) # checks, ids, spans, bucket L; a picture -> enc.patches | |
| rows = tower(enc.tile, enc.patches) # the tower graph: [G^2, 4096] picture rows (requests with a picture) | |
| x = host.row_inputs(enc, rows) # {"embeds", "valid", "cos", "sin"} at length enc.length | |
| hidden = row_graph(enc.length, x) # the row graph of that length: [1, L, 4096] after the final norm | |
| logits = host.scores(enc, hidden) # the joint head, NumPy float32: one array per question | |
| body = host.respond(request, enc, logits) # the SystemOne-compatible response | |
| or `host.decide(request, row_graph, tower)` for all of it. Refusals: errors.py (the author's ValueError conditions | |
| by name, plus the host's own: one picture at a 256 / 448 tile, at most 4,096 tokens).""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Callable | |
| import numpy as np | |
| from .errors import UnsupportedMedia | |
| from .head import JointHead | |
| from .pixels import tile_grid, tile_patches | |
| from .positions import BUCKETS, mrope_tables, padded_positions, pick_length, record_positions | |
| from .request import MAX_LENGTH, EncodedRecord, encode_record, validate_request | |
| from .response import question_probs, response | |
| from .tables import Bf16Table | |
| HIDDEN = 4096 | |
| FILES = {"tokenizer": "tokenizer.json", "embed_tokens": "embed_tokens.safetensors", "lm_head": "lm_head.safetensors", | |
| "joint_head": "joint_head.safetensors", "joint_head_config": "joint_head_config.json"} | |
| class Encoded: | |
| record: EncodedRecord | |
| length: int # the row graph length: the smallest bucket >= record.tokens | |
| tile: int | None = None # 256 | 448 for a request with a picture | |
| patches: np.ndarray | None = None # the tower graph's input [4 G^2, 1536] float32 | |
| def media_of(request: dict) -> list: | |
| """The request's one picture as a list ([] or [picture]); UnsupportedMedia for what the host does not read.""" | |
| if request.get("videos"): | |
| raise UnsupportedMedia("videos are not read by this host") | |
| if request.get("media_kwargs"): | |
| raise UnsupportedMedia("media_kwargs are not read by this host (one picture at a fixed tile)") | |
| images = list(request.get("images") or []) | |
| if len(images) > 1: | |
| raise UnsupportedMedia(f"{len(images)} pictures: this host reads one picture per request") | |
| return images | |
| class ClefHost: | |
| def __init__(self, tokenizer, embed_tokens: Bf16Table, lm_head: Bf16Table, head: JointHead, | |
| buckets=BUCKETS, max_length: int = MAX_LENGTH): | |
| if embed_tokens.shape[1] != HIDDEN or lm_head.shape[1] != HIDDEN: | |
| raise ValueError(f"tables {embed_tokens.shape} / {lm_head.shape}: expected [V, {HIDDEN}]") | |
| self.tokenizer, self.embed_tokens, self.lm_head, self.head = tokenizer, embed_tokens, lm_head, head | |
| self.buckets, self.max_length = tuple(sorted(buckets)), max_length | |
| def from_files(cls, tokenizer: str | Path, embed_tokens: str | Path, lm_head: str | Path, joint_head: str | Path, | |
| joint_head_config: str | Path, **kw) -> "ClefHost": | |
| from tokenizers import Tokenizer | |
| return cls(Tokenizer.from_file(str(tokenizer)), Bf16Table(embed_tokens), Bf16Table(lm_head), | |
| JointHead.from_files(joint_head, joint_head_config), **kw) | |
| def from_dir(cls, folder: str | Path, **kw) -> "ClefHost": | |
| folder = Path(folder) | |
| return cls.from_files(**{k: folder / v for k, v in FILES.items()}, **kw) | |
| def encode(self, request: dict, tile: int = 448) -> Encoded: | |
| """The author's checks, the host's media checks, encode_record, the bucket (RequestTooLong past 4,096).""" | |
| validate_request(request) | |
| images = media_of(request) | |
| grid = None | |
| if images: | |
| g = tile_grid(tile) | |
| grid = (g, g) | |
| record = encode_record(self.tokenizer, request, max_length=self.max_length, image_grid=grid) | |
| length = pick_length(record.tokens, self.buckets) | |
| if not images: | |
| return Encoded(record, length) | |
| return Encoded(record, length, tile, tile_patches(images[0], tile)) | |
| def row_inputs(self, enc: Encoded, image_rows: np.ndarray | None = None) -> dict[str, np.ndarray]: | |
| """The row graph's inputs at enc.length: embeds = the float32 rows of the bf16 table at the ids, the picture's | |
| rows (the tower's output, [H * W, 4096]) at the <|image_pad|> positions; valid = 1 on the record; the pad | |
| positions after it: embeds 0, valid 0; cos / sin of the record's planes (positions.py).""" | |
| rec, length = enc.record, enc.length | |
| t = rec.tokens | |
| embeds = np.zeros((1, length, HIDDEN), np.float32) | |
| embeds[0, :t] = self.embed_tokens.rows(rec.ids) | |
| if rec.image_grid is not None: | |
| h, w = rec.image_grid | |
| if image_rows is None: | |
| raise ValueError("a request with a picture needs the tower's rows") | |
| image_rows = np.asarray(image_rows, dtype=np.float32) | |
| if image_rows.shape != (h * w, HIDDEN): | |
| raise ValueError(f"picture rows {image_rows.shape}, expected {(h * w, HIDDEN)}") | |
| p = rec.token_offset | |
| embeds[0, p + 1:p + 1 + h * w] = image_rows | |
| elif image_rows is not None: | |
| raise ValueError("picture rows for a request without a picture") | |
| valid = np.zeros((1, length), np.float32) | |
| valid[0, :t] = 1.0 | |
| cos, sin = mrope_tables(padded_positions(record_positions(t, rec.token_offset, rec.image_grid), length)) | |
| return {"embeds": embeds, "valid": valid, "cos": cos, "sin": sin} | |
| def scores(self, enc: Encoded, hidden: np.ndarray, lexical: list | None = None) -> list[np.ndarray]: | |
| """The row graph's hidden ([1, L, 4096] or [T, 4096]) -> the head's logits, one float32 array per question.""" | |
| hidden = np.asarray(hidden, dtype=np.float32) | |
| if hidden.ndim == 3: | |
| hidden = hidden[0] | |
| t = enc.record.tokens | |
| return self.head.scores(hidden[:t], enc.record.ids, enc.record.questions, self.lm_head, lexical=lexical) | |
| def respond(self, request: dict, enc: Encoded, logits: list[np.ndarray]) -> dict: | |
| return response(request, enc.record.questions, [question_probs(z) for z in logits], enc.record.tokens) | |
| def decide(self, request: dict, row_graph: Callable[[int, dict], np.ndarray], | |
| tower: Callable[[int, np.ndarray], np.ndarray] | None = None, tile: int = 448) -> dict: | |
| enc = self.encode(request, tile=tile) | |
| rows = None | |
| if enc.patches is not None: | |
| if tower is None: | |
| raise UnsupportedMedia("a request with a picture needs a tower") | |
| rows = tower(enc.tile, enc.patches) | |
| hidden = row_graph(enc.length, self.row_inputs(enc, rows)) | |
| return self.respond(request, enc, self.scores(enc, hidden)) | |