src/bihmoe/tasks/bind_query.py

from __future__ import annotations
from dataclasses import dataclass, asdict
from typing import Dict, List, Tuple
import random
import json
import zlib

# Token conventions (keep compatible with current config ranges)
PAD = 0
BOS = 1
SEP = 2
NOISE0 = 3  # noise tokens are 3..(3+noise_vocab-1)

@dataclass
class BindQueryRecord:
    pairs: List[Tuple[int, int]]   # [(k,v),...]
    query_k: int
    fmt: str                       # "train" | "perturb_gap" | "braid"
    gap_max: int                   # max noise tokens inserted (gap formats)

def _rand_pairs(rng: random.Random, n_pairs: int, key_ids: List[int], val_ids: List[int]) -> List[Tuple[int,int]]:
    keys = rng.sample(key_ids, k=n_pairs)
    vals = rng.sample(val_ids, k=n_pairs)
    return list(zip(keys, vals))

def make_record(
    rng: random.Random,
    n_pairs: int,
    key_ids: List[int],
    val_ids: List[int],
    fmt: str = "train",
    gap_max: int = 0,
) -> BindQueryRecord:
    pairs = _rand_pairs(rng, n_pairs, key_ids, val_ids)
    query_k, _ = rng.choice(pairs)
    return BindQueryRecord(pairs=pairs, query_k=query_k, fmt=fmt, gap_max=gap_max)

def solve(record: BindQueryRecord) -> int:
    m = {k: v for (k, v) in record.pairs}
    return m[record.query_k]

def _stable_seed_from_record(record: BindQueryRecord) -> int:
    payload = json.dumps(asdict(record), sort_keys=True).encode("utf-8")
    return zlib.crc32(payload) & 0xffffffff

def _add_noise(rng: random.Random, ids: List[int], noise_vocab: int, gap_max: int) -> None:
    if gap_max <= 0:
        return
    gap = rng.randint(0, gap_max)
    for _ in range(gap):
        ids.append(NOISE0 + rng.randint(0, noise_vocab - 1))

def encode(
    record: BindQueryRecord,
    vocab_size: int,
    noise_vocab: int = 16,
    max_len: int = 256,
) -> Tuple[List[int], int]:
    """
    Returns:
      input_ids: list[int]
      target_id: int (value token)

    Formats:
      - train: [BOS, k1, v1, k2, v2, ..., SEP, query_k]
      - perturb_gap: [BOS, k1, noise*, v1, k2, noise*, v2, ..., SEP, query_k]
      - braid: [BOS, k1, noise*, k2, noise*, ..., SEP, v1, noise*, v2, noise*, ..., SEP, query_k]
        (Values are aligned by position with keys, not adjacency.)
    """
    rng = random.Random(_stable_seed_from_record(record))

    ids: List[int] = [BOS]

    if record.fmt in ("train", "perturb_gap"):
        for (k, v) in record.pairs:
            ids.append(k)
            if record.fmt == "perturb_gap":
                _add_noise(rng, ids, noise_vocab, record.gap_max)
            ids.append(v)
        ids.append(SEP)
        ids.append(record.query_k)

    elif record.fmt == "braid":
        # keys block
        for (k, _v) in record.pairs:
            ids.append(k)
            _add_noise(rng, ids, noise_vocab, record.gap_max)
        ids.append(SEP)
        # values block (aligned by index)
        for (_k, v) in record.pairs:
            ids.append(v)
            _add_noise(rng, ids, noise_vocab, record.gap_max)
        ids.append(SEP)
        ids.append(record.query_k)

    else:
        raise ValueError(f"Unknown fmt={record.fmt}")

    # Ensure within vocab
    for t in ids:
        if t < 0 or t >= vocab_size:
            raise ValueError(f"token {t} out of vocab_size {vocab_size}")

    target = solve(record)
    if target < 0 or target >= vocab_size:
        raise ValueError(f"target {target} out of vocab_size {vocab_size}")

    if len(ids) > max_len:
        ids = ids[:max_len]
    return ids, target

def record_to_json(record: BindQueryRecord) -> Dict:
    return asdict(record)