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, *,
        family_weight: float = 0.0, exact_family_index: Tensor | None = None,
    ) -> None:
        super().__init__()
        self.model = model
        self.adapter = adapter if adapter is not None else nn.Identity()
        self.family_weight = float(family_weight)
        if not 0.0 <= self.family_weight <= 1.0:
            raise ValueError("online family fusion weight는 0~1 범위여야 합니다.")
        if self.family_weight and exact_family_index is None:
            raise ValueError("family fusion에는 exact_family_index가 필요합니다.")
        self.register_buffer(
            "exact_family_index",
            exact_family_index if exact_family_index is not None else torch.empty(0, dtype=torch.long),
        )

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

        exact, family = self.model.forward_online(self.adapter(sequence))
        if self.family_weight:
            exact = (
                exact.log_softmax(dim=-1)
                + self.family_weight
                * family.log_softmax(dim=-1)[:, self.exact_family_index]
            )
        return exact, family


class PFormulaStudentExportWrapper06(nn.Module):
    """필요 변수: 0.6 모델·online adapter·증류 formula adapter. 작동 원리: P 수식용 두 adapter를 순서대로 합성한다."""

    def __init__(
        self,
        model: MathInk06Model,
        online_adapter: nn.Module,
        formula_adapter: nn.Module,
        *,
        family_weight: float = 0.0,
        exact_family_index: Tensor | None = None,
    ) -> None:
        super().__init__()
        self.model = model
        self.online_adapter = online_adapter
        self.formula_adapter = formula_adapter
        self.family_weight = float(family_weight)
        if not 0.0 <= self.family_weight <= 1.0:
            raise ValueError("formula family fusion weight는 0~1 범위여야 합니다.")
        if self.family_weight and exact_family_index is None:
            raise ValueError("formula family fusion에는 exact_family_index가 필요합니다.")
        self.register_buffer(
            "exact_family_index",
            exact_family_index if exact_family_index is not None
            else torch.empty(0, dtype=torch.long),
        )

    def forward(self, sequence: Tensor) -> tuple[Tensor, Tensor]:
        """필요 변수: B×128×19 formula-relative trajectory. 작동 원리: online 보정 뒤 student formula 보정을 적용해 두 logit을 반환한다."""

        adapted = self.formula_adapter(self.online_adapter(sequence))
        exact, family = self.model.classify_trajectory(adapted)
        if self.family_weight:
            exact = (
                exact.log_softmax(dim=-1)
                + self.family_weight
                * family.log_softmax(dim=-1)[:, self.exact_family_index]
            )
        return exact, family


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


class RasterDebugExportWrapper06(RasterExportWrapper06):
    """필요 변수: raster model·adapter·fusion. 작동 원리: logits와 top-4 가상 stroke 검증 출력을 함께 고정한다."""

    def forward(
        self,
        raster: Tensor,
    ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        """필요 변수: B×1×128×128 raster. 작동 원리: direct shortcut 없이 분류하고 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, coordinates, states, progress, hypothesis_scores


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,
    }