File size: 7,369 Bytes
2b605e8 | 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 | """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"}
@dataclass(frozen=True)
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
@classmethod
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)
@classmethod
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))
|