src/ink06_export.py

"""Math Ink 0.6의 online/raster 경로를 torch.export와 LiteRT 친화 출력으로 고정한다."""

from __future__ import annotations

from typing import Iterable

import torch
from torch import Tensor, nn

from .math_ink_06 import MathInk06Model, fuse_raster_logits06, virtual_features06


class OnlineExportWrapper06(nn.Module):
    """필요 변수: 0.6 모델·online adapter. 작동 원리: 실제 composite 경로의 exact/family logits를 반환한다."""

    def __init__(self, model: MathInk06Model, adapter: nn.Module | None = None) -> None:
        super().__init__()
        self.model = model
        self.adapter = adapter if adapter is not None else nn.Identity()

    def forward(self, sequence: Tensor) -> tuple[Tensor, Tensor]:
        """필요 변수: B×128×19 canonical trajectory. 작동 원리: shared encoder의 두 분류 head를 직접 실행한다."""

        return self.model.forward_online(self.adapter(sequence))


class RasterExportWrapper06(nn.Module):
    """필요 변수: 0.6 모델·raster adapter·fusion 상수. 작동 원리: top-4를 composite trajectory 경로로 분류한다."""

    def __init__(
        self, model: MathInk06Model, *, adapter: nn.Module | None = None,
        fusion_mode: str, score_weight: float,
    ) -> None:
        super().__init__()
        self.model = model
        self.adapter = adapter if adapter is not None else nn.Identity()
        self.fusion_mode = fusion_mode
        self.score_weight = float(score_weight)

    def forward(self, raster: Tensor) -> Tensor:
        """필요 변수: B×1×128×128 raster. 작동 원리: direct raster-label shortcut 없이 shared trajectory 분류를 결합한다."""

        coordinates, states, progress, hypothesis_scores = self.model.decode_raster_trajectories(raster)
        features = virtual_features06(
            coordinates, states,
            None if self.model.raster_architecture == "spatial_flat_v1" else progress,
            contract=self.model.virtual_contract,
        )
        batch, hypotheses, steps, channels = features.shape
        if self.model.use_virtual_adapter:
            raw_features = features
            internal = self.model.virtual_adapter(
                features.reshape(batch * hypotheses, steps, channels),
            ).reshape(batch, hypotheses, steps, channels)
            features = raw_features + self.model.virtual_adapter_weight * (internal - raw_features)
        flat_features = self.adapter(features.reshape(batch * hypotheses, steps, channels))
        exact, family = self.model.classify_trajectory(flat_features)
        output = {
            "hypothesis_scores": hypothesis_scores,
            "exact_logits": exact.reshape(batch, hypotheses, -1),
            "family_logits": family.reshape(batch, hypotheses, -1),
        }
        fused, _selected = fuse_raster_logits06(
            output, mode=self.fusion_mode, score_weight=self.score_weight,
        )
        return fused


def exported_equivalence06(
    eager: nn.Module, exported: torch.export.ExportedProgram, inputs: Iterable[tuple[Tensor, ...]],
) -> dict[str, float | int | bool]:
    """필요 변수: eager/export 모델·대표 입력. 작동 원리: 모든 출력 tensor의 top-1 일치와 최대 logit 오차를 계산한다."""

    exported_module = exported.module()
    samples = top1_matches = 0
    max_error = 0.0
    eager.eval()
    with torch.inference_mode():
        for arguments in inputs:
            eager_output = eager(*arguments)
            export_output = exported_module(*arguments)
            eager_values = eager_output if isinstance(eager_output, tuple) else (eager_output,)
            export_values = export_output if isinstance(export_output, tuple) else (export_output,)
            if len(eager_values) != len(export_values):
                raise ValueError("eager/export 출력 개수가 다릅니다.")
            for eager_value, export_value in zip(eager_values, export_values):
                max_error = max(max_error, float((eager_value - export_value).abs().max()))
            samples += int(eager_values[0].shape[0])
            top1_matches += int((eager_values[0].argmax(dim=-1) == export_values[0].argmax(dim=-1)).sum())
    return {
        "samples": samples, "top1_matches": top1_matches,
        "top1_agreement": top1_matches / max(samples, 1), "max_absolute_logit_error": max_error,
        "gate_passed": top1_matches == samples and max_error <= 0.02,
    }