| """Official interface for the compact auto-routed U5d candidate.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Sequence |
|
|
| import torch |
|
|
| from modchallenge.interface.base_model import ModularMultiplicationModel |
| from src.adaptive_algorithm_router_u3 import RawInputAlgorithmRouter |
| from src.compact_auto_routed_short_u5d import ( |
| CompactAutoRoutedShortU5d, |
| CompactSparseUniversalShortU5d, |
| ) |
|
|
|
|
| MAXIMUM_WIDTH = 512 |
|
|
|
|
| def _hex_digits(value: str) -> list[int]: |
| number = int(value) |
| if number < 0: |
| raise ValueError("challenge arguments must be non-negative") |
| width = max(1, (number.bit_length() + 3) // 4) |
| return [ |
| (number >> (4 * position)) & 15 |
| for position in range(width - 1, -1, -1) |
| ] |
|
|
|
|
| class CompactAutoRoutedShortU5dSubmission(ModularMultiplicationModel): |
| """Learned raw routing into compact Montgomery and short-Horner programs.""" |
|
|
| def __init__(self) -> None: |
| self.model: CompactAutoRoutedShortU5d | None = None |
| self.device: torch.device | None = None |
|
|
| def load(self, model_dir: str) -> None: |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| if self.device.type == "cpu": |
| torch.set_num_threads(4) |
| payload = torch.load( |
| Path(model_dir) / "weights.pt", |
| map_location=self.device, |
| weights_only=True, |
| ) |
| executor = CompactSparseUniversalShortU5d( |
| **payload["model_config"]["executor"] |
| ).to(self.device) |
| router = RawInputAlgorithmRouter( |
| **payload["model_config"]["router"] |
| ).to(self.device) |
| self.model = CompactAutoRoutedShortU5d( |
| executor=executor, router=router |
| ).to(self.device) |
| result = self.model.load_state_dict(payload["state_dict"], strict=True) |
| if result.missing_keys or result.unexpected_keys: |
| raise RuntimeError("compact U5d checkpoint did not load strictly") |
| self.model.eval() |
|
|
| def preprocess_a(self, a: str) -> list[int]: |
| return _hex_digits(a) |
|
|
| def preprocess_b(self, b: str) -> list[int]: |
| return _hex_digits(b) |
|
|
| def preprocess_p(self, p: str) -> list[int]: |
| return _hex_digits(p) |
|
|
| @torch.inference_mode() |
| def predict_digits(self, a_enc, b_enc, p_enc) -> list[int]: |
| return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0] |
|
|
| @torch.inference_mode() |
| def predict_digits_batch(self, inputs: Sequence[tuple]) -> list[list[int]]: |
| if self.model is None: |
| raise RuntimeError("model has not been loaded") |
| answers: list[list[int] | None] = [None] * len(inputs) |
| supported = [ |
| index |
| for index, (_a, _b, modulus) in enumerate(inputs) |
| if len(modulus) <= MAXIMUM_WIDTH |
| ] |
| if supported: |
| prediction = self.model.predict_encoded( |
| [inputs[index][0] for index in supported], |
| [inputs[index][1] for index in supported], |
| [inputs[index][2] for index in supported], |
| ) |
| for index, row in zip(supported, prediction.answers.tolist()): |
| answers[index] = [int(digit) for digit in row] |
| for index, answer in enumerate(answers): |
| if answer is None: |
| answers[index] = [0] |
| return [answer for answer in answers if answer is not None] |
|
|
| def max_batch_size(self) -> int: |
| return 20 |
|
|
|
|
| __all__ = ["CompactAutoRoutedShortU5dSubmission"] |
|
|
