File size: 3,186 Bytes
2b4770a d03104e 99431dc d03104e 2b4770a d03104e b5b2ffa 99431dc b5b2ffa 99431dc b5b2ffa 99431dc b5b2ffa 99431dc d03104e 99431dc 2b4770a 99431dc 2b4770a d03104e 99431dc b5b2ffa 99431dc d03104e 2b4770a b5b2ffa 2b4770a 99431dc b5b2ffa 99431dc b5b2ffa 99431dc b5b2ffa 99431dc d03104e 99431dc b5b2ffa d03104e b5b2ffa 2b4770a | 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 | """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"]
|