Product · Sum

What is the input width divided by the number of digits?

Why are four independent target values unnecessary if the task asks only for their sum and product?

For each training `(sum, product)` pair, enumerate the digit tuples that could have produced it.

Train per-position digit probabilities using the total probability mass of compatible tuples, then choose the most likely aggregate pair at inference.

Recognize four digit distributions and sum the probability of every digit tuple that matches the observed sum/product pair.

"""Weakly supervised digit model for NOAI China 2026 Task 3.

The model learns from the released aggregate labels (sum and product) by
marginalising over every four-digit sequence consistent with those labels.
The organizer's images and labels are required at run time and are not stored
in this repository.
"""

from __future__ import annotations

import argparse
import csv
import os
from pathlib import Path
import zipfile

import numpy as np
from PIL import Image
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset


def make_tuple_lookup() -> tuple[torch.Tensor, torch.Tensor, np.ndarray, np.ndarray]:
    """Return every ordered digit tuple and its corresponding output pair."""
    tuples = np.indices((10, 10, 10, 10), dtype=np.int64).reshape(4, -1).T
    sums = tuples.sum(axis=1)
    products = tuples.prod(axis=1)
    pairs = np.column_stack([sums, products])
    unique_pairs, pair_ids = np.unique(pairs, axis=0, return_inverse=True)
    return (
        torch.as_tensor(tuples, dtype=torch.long),
        torch.as_tensor(pair_ids, dtype=torch.long),
        unique_pairs[:, 0].astype(np.int64),
        unique_pairs[:, 1].astype(np.int64),
    )


TUPLES, TUPLE_PAIR_IDS, PAIR_SUMS, PAIR_PRODUCTS = make_tuple_lookup()
PAIR_TO_ID = {(int(s), int(p)): i for i, (s, p) in enumerate(zip(PAIR_SUMS, PAIR_PRODUCTS))}


class ProductSumDataset(Dataset):
    def __init__(self, records: list[dict], image_dir: str | Path, labeled: bool = True):
        self.records = records
        self.image_dir = Path(image_dir)
        self.labeled = labeled

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, index: int):
        record = self.records[index]
        image_path = self.image_dir / f"{record['id']}.png"
        with Image.open(image_path) as source:
            image = np.asarray(source.convert("L"), dtype=np.float32) / 255.0
        if image.shape != (28, 112):
            raise ValueError(f"Expected a 28x112 grayscale image, got {image.shape}: {image_path}")
        image = (image - 0.1307) / 0.3081
        tensor = torch.from_numpy(image).unsqueeze(0)
        if self.labeled:
            return tensor, int(record["sum"]), int(record["product"])
        return tensor, str(record["id"])


class DigitEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 96, kernel_size=3, padding=1), nn.BatchNorm2d(96), nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Sequential(nn.Flatten(), nn.Linear(96, 64), nn.ReLU(), nn.Dropout(0.15), nn.Linear(64, 10))

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        # Shared digit recognition across the four known 28-pixel segments.
        batch = images.shape[0]
        crops = torch.stack([images[..., i * 28:(i + 1) * 28] for i in range(4)], dim=1)
        crops = crops.reshape(batch * 4, 1, 28, 28)
        return self.classifier(self.features(crops)).reshape(batch, 4, 10)


def pair_probabilities(logits: torch.Tensor) -> torch.Tensor:
    log_digit_p = torch.log_softmax(logits, dim=-1)
    tuple_ids = TUPLES.to(logits.device)
    joint_log = torch.zeros(logits.shape[0], tuple_ids.shape[0], device=logits.device)
    for position in range(4):
        joint_log += log_digit_p[:, position, :][:, tuple_ids[:, position]]
    probabilities = torch.exp(joint_log)
    pair_ids = TUPLE_PAIR_IDS.to(logits.device).expand(logits.shape[0], -1)
    output = torch.zeros(logits.shape[0], len(PAIR_SUMS), device=logits.device)
    return output.scatter_add_(1, pair_ids, probabilities)


def marginal_pair_loss(logits: torch.Tensor, sums: torch.Tensor, products: torch.Tensor) -> torch.Tensor:
    pair_p = pair_probabilities(logits)
    pair_ids = [PAIR_TO_ID.get((int(s), int(p))) for s, p in zip(sums.tolist(), products.tolist())]
    if any(i is None for i in pair_ids):
        raise ValueError("A training target is not attainable by four decimal digits")
    target = torch.as_tensor(pair_ids, dtype=torch.long, device=logits.device)
    selected = pair_p.gather(1, target[:, None]).squeeze(1).clamp_min(1e-12)
    return -torch.log(selected).mean()


def read_labels(path: str | Path) -> list[dict]:
    with Path(path).open("r", newline="", encoding="utf-8-sig") as stream:
        records = list(csv.DictReader(stream))
    for row in records:
        row["sum"] = int(row["sum"])
        row["product"] = int(row["product"])
    if not records:
        raise ValueError(f"No labeled examples in {path}")
    return records


def fit_model(records: list[dict], image_dir: str | Path, device: torch.device, epochs: int = 25, seed: int = 42) -> DigitEncoder:
    torch.manual_seed(seed)
    np.random.seed(seed)
    loader = DataLoader(ProductSumDataset(records, image_dir), batch_size=96, shuffle=True, num_workers=2, pin_memory=device.type == "cuda")
    model = DigitEncoder().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1.5e-3, weight_decay=2e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    for epoch in range(epochs):
        model.train()
        total, seen = 0.0, 0
        for images, sums, products in loader:
            images = images.to(device, non_blocking=True)
            sums, products = sums.to(device), products.to(device)
            loss = marginal_pair_loss(model(images), sums, products)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            total += float(loss.detach()) * len(images)
            seen += len(images)
        scheduler.step()
        print(f"epoch {epoch + 1:02d}/{epochs}: aggregate negative log likelihood={total / seen:.4f}")
    return model


@torch.no_grad()
def predict(
    model: DigitEncoder,
    image_dir: str | Path,
    device: torch.device,
    batch_size: int = 128,
    sample_ids: set[str] | None = None,
) -> list[dict]:
    paths = sorted(Path(image_dir).glob("*.png"))
    if sample_ids is not None:
        paths = [p for p in paths if p.stem in sample_ids]
    if not paths:
        raise FileNotFoundError(f"No PNG images in {image_dir}")
    records = [{"id": p.stem} for p in paths]
    loader = DataLoader(ProductSumDataset(records, image_dir, labeled=False), batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=device.type == "cuda")
    model.eval()
    rows: list[dict] = []
    for images, sample_ids in loader:
        image_logits = model(images.to(device, non_blocking=True))
        pair_p = pair_probabilities(image_logits).cpu().numpy()
        best = pair_p.argmax(axis=1)
        for sample_id, pair_id in zip(sample_ids, best):
            rows.append({"id": str(sample_id), "sum": int(PAIR_SUMS[pair_id]), "product": int(PAIR_PRODUCTS[pair_id])})
    return rows


def write_submission(records: list[dict], path: str | Path) -> None:
    with Path(path).open("w", newline="", encoding="utf-8") as stream:
        writer = csv.DictWriter(stream, fieldnames=["id", "sum", "product"])
        writer.writeheader()
        writer.writerows(records)


def accuracy(records: list[dict], expected: list[dict]) -> tuple[float, float, float]:
    pred = {str(r["id"]): r for r in records}
    gold = {str(r["id"]): r for r in expected}
    common = sorted(set(pred) & set(gold))
    if len(common) != len(gold):
        raise ValueError("Validation predictions do not cover every held-out id")
    sum_acc = np.mean([pred[i]["sum"] == gold[i]["sum"] for i in common])
    product_acc = np.mean([pred[i]["product"] == gold[i]["product"] for i in common])
    return float((sum_acc + product_acc) / 2), float(sum_acc), float(product_acc)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--mode", choices=("validate", "submit"), default="submit")
    parser.add_argument("--train-dir", default="/bohr/train-rppd/v1")
    parser.add_argument("--data-dir", default=os.environ.get("DATA_PATH", "/bohr"))
    parser.add_argument("--output-dir", default=".")
    parser.add_argument("--epochs", type=int, default=25)
    args = parser.parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_dir = Path(args.train_dir)
    labels_path = next((p for p in (train_dir / "train_labels.csv", train_dir / "train.csv") if p.is_file()), None)
    if labels_path is None:
        raise FileNotFoundError(f"Expected train_labels.csv or train.csv under {train_dir}")
    labels = read_labels(labels_path)
    if args.mode == "validate":
        rng = np.random.default_rng(42)
        order = rng.permutation(len(labels))
        boundary = int(0.9 * len(order))
        train_records = [labels[i] for i in order[:boundary]]
        val_records = [labels[i] for i in order[boundary:]]
        model = fit_model(train_records, train_dir / "train_images", device, args.epochs)
        val_ids = {str(r["id"]) for r in val_records}
        val_preds = predict(model, train_dir / "train_images", device, sample_ids=val_ids)
        _, sum_acc, product_acc = accuracy(val_preds, val_records)
        print(f"holdout exact accuracy: sum={sum_acc:.6f}, product={product_acc:.6f}")
        return

    model = fit_model(labels, train_dir / "train_images", device, args.epochs)
    data_dir = Path(args.data_dir)
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    for split in ("val", "test"):
        rows = predict(model, data_dir / split, device)
        write_submission(rows, out_dir / f"submission_{split}.csv")
    with zipfile.ZipFile(out_dir / "submission.zip", "w", compression=zipfile.ZIP_DEFLATED) as archive:
        archive.write(out_dir / "submission_val.csv", "submission_val.csv")
        archive.write(out_dir / "submission_test.csv", "submission_test.csv")
    print(f"Wrote {out_dir / 'submission.zip'} on {device}")


if __name__ == "__main__":
    main()

Download solution.py