encoding.py

"""Shared sequence encoding for the modmul BP-install model.

The model learns the field-multiplication map  (p, x, y) -> (x*y) mod p  with
x, y in [0, p).  At inference the legal two-operand reduction a%p, b%p produces
x, y (the same step the baselines use); training samples x, y in [0, p) directly,
which matches the (a mod p) distribution for p << operand range.

Fixed-width, base-10, reverse-LSB answer (the proven arithmetic recipe). The
prime p is an explicit in-sequence conditioning field so one model generalises
across primes:

    p007*012,003=R120\n        # W=3 example: p=7, x=12%7=5? no -- x,y already <p
                               # e.g. p=7,x=5,y=4 -> 20%7=6 -> "p7*5,4=R6\n" (W=1)

Layout (width W = max decimal digits of any prime in scope):
    "p" P(W) "*" X(W) "," Y(W) "=R" REV_ANS(W) "\n"
zero-padded fields; REV_ANS is the W-digit answer reversed (least-significant
digit first), so digit k is emitted at a fixed position and the carry flows
left-to-right the way reverse-LSB addition does.
"""
from __future__ import annotations

MUL, SEP, EQ, REVMARK, NL = "*", ",", "=", "R", "\n"
NEWLINE_ID = ord("\n")


def width_for_primes(primes) -> int:
    """Decimal digits needed to hold the largest residue (p-1)."""
    return max(len(str(p - 1)) for p in primes)


def prompt_str(p: int, x: int, y: int, W: int) -> str:
    return f"p{p:0{W}d}*{x:0{W}d},{y:0{W}d}={REVMARK}"


def build_example(p: int, x: int, y: int, W: int):
    """Return (text, ann) where ann[t] holds the install label for the residual
    at byte position t (which predicts byte t+1), NTP-aligned."""
    ans = (x * y) % p
    rev_ans = f"{ans:0{W}d}"[::-1]            # LSB-first, width W
    prompt = prompt_str(p, x, y, W)
    text = prompt + rev_ans + NL
    ann = [dict() for _ in range(len(text))]
    base = len(prompt)                        # first answer char index
    for k in range(W):                        # answer digit k (LSB-first)
        pos = base + k
        ann[pos - 1]["ans"] = int(rev_ans[k])  # residual at pos-1 predicts it
    return text, ann


def answer_len(W: int) -> int:
    return W


# --- Scratchpad ("chain-of-thought") inference variant -------------------
# The model GENERATES the product as reverse-LSB digits before the answer:
#     "p" P(W) "*" X(W) "," Y(W) "=" REV_PROD(2W) "R" REV_ANS(W) "\n"
# At inference we feed up to '=' and greedily decode the full scratchpad; the
# answer is the last W generated digits. Python never computes x*y -- the model
# emits the product digits (a learned circuit). Only a%p, b%p is done in code.

def prompt_str_sp(p: int, x: int, y: int, W: int) -> str:
    return f"p{p:0{W}d}*{x:0{W}d},{y:0{W}d}{EQ}"


def scratchpad_len(W: int) -> int:
    return 2 * W + 1 + W


def decode_answer_sp(gen_chars: str, W: int) -> int:
    """Generated suffix is 2W product digits, 'R', then W answer digits."""
    return decode_answer(gen_chars[2 * W + 1 : 2 * W + 1 + W], W)


# --- Schoolbook ("partial products") inference variant -------------------
# Generated suffix: W partial products (W+1 wide), 'P', 2W product digits,
# 'R', W answer digits. The model emits the partial products and sums them
# (learned circuit); the answer is the final W digits.

def school_gen_len(W: int) -> int:
    return W * (W + 1) + 1 + 2 * W + 1 + W


def decode_answer_school(gen_chars: str, W: int) -> int:
    start = W * (W + 1) + 1 + 2 * W + 1
    return decode_answer(gen_chars[start:start + W], W)


def decode_answer(gen_chars: str, W: int) -> int:
    """Greedy-generated W chars are the reverse-LSB answer digits. Non-digits
    (untrained junk) count as 0. Returns the integer answer."""
    val = 0
    for k in range(min(W, len(gen_chars))):
        c = gen_chars[k]
        d = int(c) if "0" <= c <= "9" else 0  # ASCII only ('³'.isdigit() is True but int() fails)
        val += d * (10 ** k)
    return val