| """Official interface for compiled-local Montgomery v42.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Sequence |
|
|
| import torch |
|
|
| from modchallenge.interface.base_model import ModularMultiplicationModel |
| from src.compact_sparse_montgomery_v39 import CompiledLocalSparseMontgomeryV42 |
|
|
|
|
| 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 CompiledLocalMontgomeryV42Submission(ModularMultiplicationModel): |
| """Raw arguments through v20 arithmetic and compiled learned local maps.""" |
|
|
| def __init__(self) -> None: |
| self.model: CompiledLocalSparseMontgomeryV42 | 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, |
| ) |
| self.model = CompiledLocalSparseMontgomeryV42( |
| **payload["model_config"] |
| ).to(self.device) |
| result = self.model.load_state_dict(payload["state_dict"], strict=True) |
| if result.missing_keys or result.unexpected_keys: |
| raise RuntimeError("v42 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__ = ["CompiledLocalMontgomeryV42Submission"] |
|
|
