| """Complete-input 1K product profile around an unchanged native runtime.""" |
| from dataclasses import replace |
|
|
| MAX_INPUT_TOKENS = 1024 |
|
|
|
|
| class Complete1KCollator: |
| """Use the native collator's actual length error, never truncate fields. |
| |
| Constructible with a tokenizer before any model is loaded. A whole request |
| is admitted before predict, and every physical batch before tensor assembly. |
| """ |
| def __init__(self, native_collator): |
| self.base = type(native_collator)(native_collator.tokenizer, |
| max_length=MAX_INPUT_TOKENS, |
| state_truncation="error") |
| self.tokenizer = self.base.tokenizer |
| self.marker, self.pad = self.base.marker, self.base.pad |
| self.tensor_batches = 0 |
|
|
| def tokens(self, text): |
| return self.base.tokens(text) |
|
|
| def encode(self, row, labeled=False): |
| encoded = self.base.encode(row, labeled=labeled) |
| if (encoded["input_tokens"] > MAX_INPUT_TOKENS |
| or encoded["state_tokens_original"] != encoded["state_tokens_kept"]): |
| raise ValueError("Native collator violated the complete-input 1K profile") |
| return encoded |
|
|
| def admit(self, records): |
| return [self.encode(row, labeled=False) for row in records] |
|
|
| def __call__(self, records, labeled=False, device="cpu"): |
| records = list(records) |
| |
| for row in records: |
| self.encode(row, labeled=labeled) |
| self.tensor_batches += 1 |
| return self.base(records, labeled=labeled, device=device) |
|
|
|
|
|
|
| def predict_1k(native, records, *, batch_size=8): |
| """Admit every complete input, then reuse its encoding within this request. |
| |
| Native prediction, batch boundaries and outputs are unchanged. Encodings |
| are released with the synchronous call; no state activations are cached. |
| """ |
| from ._request import predict_1k as predict_admitted |
| return predict_admitted(native, records, batch_size=batch_size) |
|
|
