| """AIFlow Math Ink 0.6의 shared trajectory와 raster→virtual-stroke 신경망이다.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from pathlib import Path |
| import time |
| from typing import Any, Sequence |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
| from torch import Tensor, nn |
|
|
| from .ink06_canonical import FEATURE_NAMES_06, MAX_EVENTS, canonicalize_ink06 |
|
|
|
|
| class ResidualTcnBlock06(nn.Module): |
| """필요 변수: channel·dilation. 작동 원리: 모바일 호환 Conv1d residual로 타점 패턴을 인코딩한다.""" |
|
|
| def __init__(self, channels: int, dilation: int) -> None: |
| super().__init__() |
| groups = 8 if channels % 8 == 0 else 1 |
| self.network = nn.Sequential( |
| nn.Conv1d(channels, channels, 5, padding=2 * dilation, dilation=dilation), |
| nn.GroupNorm(groups, channels), nn.GELU(), nn.Dropout(0.10), |
| nn.Conv1d(channels, channels, 1), nn.GroupNorm(groups, channels), |
| ) |
| self.activation = nn.GELU() |
|
|
| def forward(self, value: Tensor) -> Tensor: |
| """필요 변수: B×C×T. 작동 원리: 동일 길이 residual feature를 반환한다.""" |
|
|
| return self.activation(value + self.network(value)) |
|
|
|
|
| class SharedTrajectoryEncoder06(nn.Module): |
| """필요 변수: 19채널·hidden. 작동 원리: padding을 제외한 attention 통계로 shared embedding을 만든다.""" |
|
|
| def __init__(self, input_size: int = len(FEATURE_NAMES_06), hidden_size: int = 128) -> None: |
| super().__init__() |
| groups = 8 if hidden_size % 8 == 0 else 1 |
| self.input_projection = nn.Sequential( |
| nn.Conv1d(input_size, hidden_size, 1), nn.GroupNorm(groups, hidden_size), nn.GELU(), |
| ) |
| self.blocks = nn.Sequential(*(ResidualTcnBlock06(hidden_size, dilation) for dilation in (1, 2, 4, 8))) |
| self.attention = nn.Conv1d(hidden_size, 1, 1) |
|
|
| def forward(self, sequence: Tensor) -> Tensor: |
| """필요 변수: B×128×19. 작동 원리: stroke_progress=-1 padding을 attention/통계에서 제거한다.""" |
|
|
| mask = sequence[:, :, 8] >= 0 |
| encoded = self.blocks(self.input_projection(sequence.transpose(1, 2))) |
| attention = self.attention(encoded).masked_fill(~mask.unsqueeze(1), -1e4) |
| weights = attention.softmax(dim=2) |
| mean = (encoded * weights).sum(dim=2) |
| variance = ((encoded - mean.unsqueeze(2)).square() * weights).sum(dim=2) |
| maximum = encoded.masked_fill(~mask.unsqueeze(1), -1e4).amax(dim=2) |
| return torch.cat((mean, maximum, torch.sqrt(variance.clamp_min(1e-6))), dim=1) |
|
|
|
|
| class VirtualTrajectoryAdapter06(nn.Module): |
| """필요 변수: virtual 19채널 feature. 작동 원리: 온라인 계약 채널을 보존하며 raster 전용 residual 보정을 학습한다.""" |
|
|
| def __init__(self, channels: int = len(FEATURE_NAMES_06), hidden_size: int = 48) -> None: |
| super().__init__() |
| self.network = nn.Sequential( |
| nn.Conv1d(channels, hidden_size, 1), nn.GELU(), |
| nn.Conv1d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size), nn.GELU(), |
| nn.Conv1d(hidden_size, channels, 1), |
| ) |
| nn.init.zeros_(self.network[-1].weight) |
| nn.init.zeros_(self.network[-1].bias) |
| |
| mutable = torch.ones(channels) |
| mutable[[7, 8, 17, 18]] = 0.0 |
| self.register_buffer("mutable_channels", mutable.view(1, 1, channels), persistent=False) |
|
|
| def forward(self, sequence: Tensor) -> Tensor: |
| """필요 변수: B×128×19 feature. 작동 원리: zero-init residual을 허용 채널에만 더한다.""" |
|
|
| delta = self.network(sequence.transpose(1, 2)).transpose(1, 2) |
| return sequence + delta * self.mutable_channels |
|
|
|
|
| class DepthwiseRasterEncoder06(nn.Module): |
| """필요 변수: 128×128 grayscale. 작동 원리: depthwise CNN의 8×8 공간 배치를 보존해 vectorizer에 전달한다.""" |
|
|
| def __init__(self, hidden_size: int) -> None: |
| super().__init__() |
| channels = (16, 32, 64, hidden_size) |
| layers: list[nn.Module] = [nn.Conv2d(1, channels[0], 3, stride=2, padding=1), nn.GELU()] |
| for source, target in zip(channels, channels[1:]): |
| layers.extend([ |
| nn.Conv2d(source, source, 3, stride=2, padding=1, groups=source), |
| nn.Conv2d(source, target, 1), nn.GroupNorm(8 if target % 8 == 0 else 1, target), nn.GELU(), |
| ]) |
| self.network = nn.Sequential(*layers) |
| self.spatial_projection = nn.Sequential( |
| nn.Flatten(), nn.Linear(hidden_size * 8 * 8, hidden_size), nn.LayerNorm(hidden_size), nn.GELU(), |
| ) |
| self.position_projection = nn.Linear(2, hidden_size, bias=False) |
| axis = torch.linspace(-1.0, 1.0, 8) |
| grid_y, grid_x = torch.meshgrid(axis, axis, indexing="ij") |
| self.register_buffer("spatial_positions", torch.stack((grid_x, grid_y), dim=-1).view(64, 2), persistent=False) |
| self.fine_projection = nn.Conv2d(64, hidden_size, 1) |
| fine_axis = torch.linspace(-1.0, 1.0, 16) |
| fine_y, fine_x = torch.meshgrid(fine_axis, fine_axis, indexing="ij") |
| self.register_buffer( |
| "fine_positions", torch.stack((fine_x, fine_y), dim=-1).view(256, 2), persistent=False, |
| ) |
| self.pointer_projection = nn.Conv2d(32, hidden_size, 1) |
| pointer_axis = (torch.arange(32, dtype=torch.float32) + 0.5) / 32.0 |
| pointer_y, pointer_x = torch.meshgrid(pointer_axis, pointer_axis, indexing="ij") |
| self.register_buffer( |
| "pointer_positions", torch.stack((pointer_x, pointer_y), dim=-1).view(1024, 2), persistent=False, |
| ) |
|
|
| def forward( |
| self, raster: Tensor, *, fine_tokens: bool = False, pointer_tokens: bool = False, |
| ) -> tuple[Tensor, Tensor]: |
| """필요 변수: B×1×128×128·해상도 선택. 작동 원리: 전역 요약과 8/16/32-grid 위치 token을 반환한다.""" |
|
|
| feature = raster |
| fine_feature = None |
| pointer_feature = None |
| for index, layer in enumerate(self.network): |
| feature = layer(feature) |
| if index == 5: |
| pointer_feature = feature |
| if index == 9: |
| fine_feature = feature |
| if pointer_tokens: |
| if pointer_feature is None: |
| raise RuntimeError("32×32 raster feature가 생성되지 않았습니다.") |
| tokens = self.pointer_projection(pointer_feature).flatten(2).transpose(1, 2) |
| positions = self.pointer_positions * 2.0 - 1.0 |
| tokens = tokens + self.position_projection(positions).unsqueeze(0) |
| elif fine_tokens: |
| if fine_feature is None: |
| raise RuntimeError("16×16 raster feature가 생성되지 않았습니다.") |
| tokens = self.fine_projection(fine_feature).flatten(2).transpose(1, 2) |
| tokens = tokens + self.position_projection(self.fine_positions).unsqueeze(0) |
| else: |
| tokens = feature.flatten(2).transpose(1, 2) |
| tokens = tokens + self.position_projection(self.spatial_positions).unsqueeze(0) |
| return self.spatial_projection(feature), tokens |
|
|
|
|
| class RasterCrossAttentionBlock06(nn.Module): |
| """필요 변수: trajectory query·4×4 raster token. 작동 원리: 각 가상 타점이 대응할 이미지 위치를 직접 조회한다.""" |
|
|
| def __init__(self, hidden_size: int) -> None: |
| super().__init__() |
| heads = 4 if hidden_size % 4 == 0 else 1 |
| self.query_norm = nn.LayerNorm(hidden_size) |
| self.memory_norm = nn.LayerNorm(hidden_size) |
| self.attention = nn.MultiheadAttention(hidden_size, heads, batch_first=True) |
| self.output_norm = nn.LayerNorm(hidden_size) |
| self.residual_gate = nn.Parameter(torch.zeros(())) |
|
|
| def forward(self, query: Tensor, memory: Tensor, *, gated: bool = False) -> Tensor: |
| """필요 변수: query·공간 memory·gate 여부. 작동 원리: 위치 증거를 직접 또는 zero-init residual로 합친다.""" |
|
|
| attended, _weights = self.attention( |
| self.query_norm(query), self.memory_norm(memory), self.memory_norm(memory), need_weights=False, |
| ) |
| if gated: |
| return query + torch.tanh(self.residual_gate) * attended |
| return self.output_norm(query + attended) |
|
|
|
|
| class VirtualStrokeDecoder06(nn.Module): |
| """필요 변수: raster embedding·가설 수. 작동 원리: 4-layer causal Conv1d가 top-k 좌표와 pen state를 만든다.""" |
|
|
| def __init__(self, hidden_size: int = 128, hypotheses: int = 4, max_events: int = MAX_EVENTS) -> None: |
| super().__init__() |
| self.hypotheses = hypotheses |
| self.max_events = max_events |
| self.query = nn.Parameter(torch.randn(max_events, hidden_size) * 0.02) |
| self.hypothesis = nn.Embedding(hypotheses, hidden_size) |
| self.decoder = nn.ModuleList([ |
| nn.Sequential( |
| nn.Conv1d(hidden_size, hidden_size, kernel_size=5), |
| nn.GroupNorm(8 if hidden_size % 8 == 0 else 1, hidden_size), nn.GELU(), |
| ) |
| for _ in range(4) |
| ]) |
| self.cross_attention = RasterCrossAttentionBlock06(hidden_size) |
| self.coordinate_head = nn.Linear(hidden_size, 2) |
| self.state_head = nn.Linear(hidden_size, 3) |
| self.progress_head = nn.Linear(hidden_size, 1) |
| self.score_head = nn.Linear(hidden_size, 1) |
| self.pointer_query = nn.Linear(hidden_size, hidden_size, bias=False) |
| self.pointer_key = nn.Linear(hidden_size, hidden_size, bias=False) |
| self.pointer_temperature = 0.5 |
| self.pointer_logits_for_loss: Tensor | None = None |
|
|
| def forward( |
| self, embedding: Tensor, spatial_tokens: Tensor | None = None, *, gated_attention: bool = False, |
| pointer_positions: Tensor | None = None, ink_prior: Tensor | None = None, |
| ) -> tuple[Tensor, Tensor, Tensor, Tensor]: |
| """필요 변수: B×H 요약·선택 spatial token/ink prior. 작동 원리: causal path와 선택적 ink-pointer로 top-4 궤적을 반환한다.""" |
|
|
| batch = embedding.shape[0] |
| self.pointer_logits_for_loss = None |
| query = self.query.view(1, 1, self.max_events, -1) |
| hypothesis = self.hypothesis.weight.view(1, self.hypotheses, 1, -1) |
| value = query + hypothesis + embedding.view(batch, 1, 1, -1) |
| value = value.reshape(batch * self.hypotheses, self.max_events, -1).transpose(1, 2) |
| memory = None |
| if spatial_tokens is not None: |
| memory = spatial_tokens.unsqueeze(1).expand(-1, self.hypotheses, -1, -1) |
| memory = memory.reshape(batch * self.hypotheses, spatial_tokens.shape[1], spatial_tokens.shape[2]) |
| |
| if pointer_positions is None: |
| value = self.cross_attention( |
| value.transpose(1, 2), memory, gated=gated_attention, |
| ).transpose(1, 2) |
| for layer in self.decoder: |
| value = value + layer(nn.functional.pad(value, (4, 0))) |
| decoded = value.transpose(1, 2) |
| if pointer_positions is not None: |
| if memory is None or ink_prior is None: |
| raise ValueError("ink pointer에는 spatial memory와 ink prior가 모두 필요합니다.") |
| if pointer_positions.shape != (memory.shape[1], 2) or ink_prior.shape != (batch, memory.shape[1]): |
| raise ValueError("ink pointer position/prior shape가 spatial token과 일치하지 않습니다.") |
| query = self.pointer_query(decoded) |
| key = self.pointer_key(memory) |
| pointer_logits = torch.bmm(query, key.transpose(1, 2)) / (decoded.shape[-1] ** 0.5) |
| expanded_prior = ink_prior[:, None].expand(-1, self.hypotheses, -1).reshape( |
| batch * self.hypotheses, memory.shape[1], |
| ) |
| |
| pointer_logits = pointer_logits + 2.5 * (expanded_prior + 1e-4).log().unsqueeze(1) |
| self.pointer_logits_for_loss = pointer_logits.view( |
| batch, self.hypotheses, self.max_events, memory.shape[1], |
| ) |
| soft_probability = (pointer_logits / self.pointer_temperature).softmax(dim=-1) |
| hard_probability = nn.functional.one_hot( |
| soft_probability.argmax(dim=-1), num_classes=soft_probability.shape[-1], |
| ).to(dtype=soft_probability.dtype) |
| |
| pointer_probability = soft_probability if self.training else hard_probability |
| coordinates = torch.matmul(pointer_probability, pointer_positions.to(decoded)).view( |
| batch, self.hypotheses, self.max_events, 2, |
| ) |
| else: |
| coordinates = self.coordinate_head(decoded).sigmoid().view(batch, self.hypotheses, self.max_events, 2) |
| states = self.state_head(decoded).view(batch, self.hypotheses, self.max_events, 3) |
| progress = self.progress_head(decoded).sigmoid().view(batch, self.hypotheses, self.max_events) |
| scores = self.score_head(decoded[:, -1]).view(batch, self.hypotheses) |
| return coordinates, states, progress, scores |
|
|
|
|
| def virtual_features06( |
| coordinates: Tensor, state_logits: Tensor, stroke_progress: Tensor | None = None, |
| *, contract: str = "legacy_v1", |
| ) -> Tensor: |
| """필요 변수: 좌표·state·progress·계약. 작동 원리: virtual stroke를 19채널 shared encoder 입력으로 변환한다.""" |
|
|
| if contract not in {"legacy_v1", "canonical_v2"}: |
| raise ValueError("지원하지 않는 virtual feature contract입니다.") |
|
|
| batch, hypotheses, steps, _axis = coordinates.shape |
| probability = state_logits.softmax(dim=-1) |
| pen_start = probability[..., 1] |
| minimum = coordinates.amin(dim=2, keepdim=True) |
| span = (coordinates.amax(dim=2, keepdim=True) - minimum).clamp_min(1e-8 if contract == "canonical_v2" else 1e-5) |
| shape = (coordinates - minimum) / span |
| canvas_delta = torch.cat((torch.zeros_like(coordinates[:, :, :1]), coordinates[:, :, 1:] - coordinates[:, :, :-1]), dim=2) |
| shape_delta = torch.cat((torch.zeros_like(shape[:, :, :1]), shape[:, :, 1:] - shape[:, :, :-1]), dim=2) |
| delta = shape_delta if contract == "canonical_v2" else canvas_delta |
| delta = delta * (1.0 - pen_start).unsqueeze(-1) |
| distance = delta.square().sum(dim=-1, keepdim=True).clamp_min(1e-8).sqrt() |
| direction = delta / distance |
| previous = torch.cat((torch.zeros_like(direction[:, :, :1]), direction[:, :, :-1]), dim=2) |
| curvature = previous[..., 0] * direction[..., 1] - previous[..., 1] * direction[..., 0] |
| progress = stroke_progress |
| if progress is None: |
| progress = torch.linspace(0.0, 1.0, steps, device=coordinates.device).view(1, 1, steps).expand(batch, hypotheses, -1) |
| aspect = (span[..., 0] / span[..., 1]).expand(-1, -1, steps) |
| ones = torch.ones_like(progress) |
| bbox_top = minimum[..., 1].expand(-1, -1, steps) |
| bbox_bottom = (minimum[..., 1] + span[..., 1]).expand(-1, -1, steps) |
| bbox_height = span[..., 1].expand(-1, -1, steps) |
| center_y = ((bbox_top + bbox_bottom) * 0.5) |
| if contract == "canonical_v2": |
| canvas_distance = canvas_delta.square().sum(dim=-1).sqrt() * 128.0 |
| time_delta = canvas_distance / (8.0 * 6.0) |
| speed = torch.where(canvas_distance > 1e-8, torch.full_like(canvas_distance, 48.0 / 256.0), torch.zeros_like(canvas_distance)) |
| else: |
| time_delta = (1.0 / (6.0 * steps)) * ones |
| speed = canvas_delta.square().sum(dim=-1).clamp_min(1e-8).sqrt() * 6.0 |
| features = torch.stack(( |
| shape[..., 0], shape[..., 1], coordinates[..., 0], coordinates[..., 1], |
| direction[..., 0], direction[..., 1], curvature, pen_start, progress, aspect, |
| bbox_top, bbox_bottom, bbox_height, center_y, ones, time_delta, |
| speed, ones, ones, |
| ), dim=-1) |
| if contract == "legacy_v1": |
| valid = 1.0 - probability[..., 2] |
| features[..., 8] = torch.where(valid > 0.5, features[..., 8], -torch.ones_like(features[..., 8])) |
| return features |
|
|
|
|
| def equivalent_trajectory_targets06( |
| coordinates: Tensor, states: Tensor, hypotheses: int = 4, |
| ) -> tuple[Tensor, Tensor, Tensor]: |
| """필요 변수: B×T 좌표·state. 작동 원리: 같은 raster를 만드는 방향/획순서 대안 trajectory를 생성한다.""" |
|
|
| if hypotheses != 4: |
| raise ValueError("현재 equivalent target 계약은 top-4 전용입니다.") |
| coordinate_batches: list[Tensor] = [] |
| state_batches: list[Tensor] = [] |
| progress_batches: list[Tensor] = [] |
| for sample_coordinates, sample_states in zip(coordinates, states, strict=True): |
| starts = torch.nonzero(sample_states == 1, as_tuple=False).flatten().tolist() |
| if not starts or starts[0] != 0: |
| starts.insert(0, 0) |
| starts = sorted(set(int(value) for value in starts if int(value) < len(sample_states))) |
| boundaries = starts + [len(sample_states)] |
| strokes = [sample_coordinates[boundaries[index]:boundaries[index + 1]] for index in range(len(starts))] |
| variants = ( |
| strokes, |
| [stroke.flip(0) for stroke in strokes], |
| list(reversed(strokes)), |
| [stroke.flip(0) for stroke in reversed(strokes)], |
| ) |
| sample_coordinate_targets = [] |
| sample_state_targets = [] |
| sample_progress_targets = [] |
| for variant in variants: |
| joined = torch.cat(variant, dim=0) |
| target_states = torch.zeros(len(joined), dtype=states.dtype, device=states.device) |
| target_progress = torch.zeros(len(joined), dtype=coordinates.dtype, device=coordinates.device) |
| cursor = 0 |
| for stroke in variant: |
| target_states[cursor] = 1 |
| target_progress[cursor:cursor + len(stroke)] = torch.linspace( |
| 0.0, 1.0, len(stroke), dtype=coordinates.dtype, device=coordinates.device, |
| ) |
| cursor += len(stroke) |
| target_states[-1] = 2 |
| sample_coordinate_targets.append(joined) |
| sample_state_targets.append(target_states) |
| sample_progress_targets.append(target_progress) |
| coordinate_batches.append(torch.stack(sample_coordinate_targets)) |
| state_batches.append(torch.stack(sample_state_targets)) |
| progress_batches.append(torch.stack(sample_progress_targets)) |
| return torch.stack(coordinate_batches), torch.stack(state_batches), torch.stack(progress_batches) |
|
|
|
|
| def equivalent_modality_features06(sequence: Tensor) -> Tensor: |
| """필요 변수: B×128×19 online feature. 작동 원리: raster 모드용 방향/획순서 불변 variant 네 개를 만든다.""" |
|
|
| batches: list[Tensor] = [] |
| for sample in sequence: |
| starts = torch.nonzero(sample[:, 7] > 0.5, as_tuple=False).flatten().tolist() |
| if not starts or starts[0] != 0: |
| starts.insert(0, 0) |
| starts = sorted(set(int(value) for value in starts if int(value) < len(sample))) |
| boundaries = starts + [len(sample)] |
| strokes = [sample[boundaries[index]:boundaries[index + 1]] for index in range(len(starts))] |
| variants = ( |
| strokes, [stroke.flip(0) for stroke in strokes], list(reversed(strokes)), |
| [stroke.flip(0) for stroke in reversed(strokes)], |
| ) |
| rows = [] |
| for variant in variants: |
| value = torch.cat(variant, dim=0).clone() |
| value[:, 7] = 0.0 |
| cursor = 0 |
| for stroke in variant: |
| value[cursor, 7] = 1.0 |
| value[cursor:cursor + len(stroke), 8] = torch.linspace( |
| 0.0, 1.0, len(stroke), device=value.device, dtype=value.dtype, |
| ) |
| cursor += len(stroke) |
| delta_shape = torch.cat((torch.zeros_like(value[:1, :2]), value[1:, :2] - value[:-1, :2]), dim=0) |
| delta_canvas = torch.cat((torch.zeros_like(value[:1, 2:4]), value[1:, 2:4] - value[:-1, 2:4]), dim=0) |
| delta_shape[value[:, 7] > 0.5] = 0.0 |
| delta_canvas[value[:, 7] > 0.5] = 0.0 |
| distance_shape = delta_shape.square().sum(dim=-1).sqrt() |
| direction = delta_shape / distance_shape.clamp_min(1e-6).unsqueeze(-1) |
| previous = torch.cat((torch.zeros_like(direction[:1]), direction[:-1]), dim=0) |
| value[:, 4:6] = direction |
| value[:, 6] = previous[:, 0] * direction[:, 1] - previous[:, 1] * direction[:, 0] |
| canvas_distance = delta_canvas.square().sum(dim=-1).sqrt() * 128.0 |
| value[:, 15] = canvas_distance / 48.0 |
| value[:, 16] = torch.where( |
| canvas_distance > 1e-8, torch.full_like(canvas_distance, 48.0 / 256.0), |
| torch.zeros_like(canvas_distance), |
| ) |
| value[:, 17] = 1.0 |
| value[:, 18] = 1.0 |
| rows.append(value) |
| batches.append(torch.stack(rows)) |
| return torch.stack(batches) |
|
|
|
|
| def soft_rasterize_virtual06( |
| coordinates: Tensor, *, size: int = 32, sigma: float = 0.025, point_stride: int = 1, |
| point_weights: Tensor | None = None, |
| ) -> Tensor: |
| """필요 변수: B×K×T 좌표·선택 weight. 작동 원리: END/padding을 제외한 대표 타점을 부드러운 raster로 변환한다.""" |
|
|
| if size <= 0 or sigma <= 0 or point_stride <= 0: |
| raise ValueError("raster size·sigma·point_stride는 양수여야 합니다.") |
| original_time_shape = coordinates.shape[:-1] |
| if point_weights is not None: |
| if point_weights.shape != original_time_shape: |
| raise ValueError("point weight는 coordinate의 원본 B×K×T 축과 일치해야 합니다.") |
| coordinates = coordinates[:, :, ::point_stride] |
| if point_weights is not None: |
| point_weights = point_weights[:, :, ::point_stride] |
| axis = (torch.arange(size, device=coordinates.device, dtype=coordinates.dtype) + 0.5) / size |
| grid_y, grid_x = torch.meshgrid(axis, axis, indexing="ij") |
| grid = torch.stack((grid_x, grid_y), dim=-1) |
| minimum = torch.full((*coordinates.shape[:2], size, size), torch.inf, dtype=coordinates.dtype, device=coordinates.device) |
| weighted_maximum = torch.zeros((*coordinates.shape[:2], size, size), dtype=coordinates.dtype, device=coordinates.device) |
| |
| offset = 0 |
| for chunk in coordinates.split(32, dim=2): |
| distance = (chunk[:, :, :, None, None] - grid).square().sum(dim=-1) |
| if point_weights is None: |
| minimum = torch.minimum(minimum, distance.amin(dim=2)) |
| else: |
| weights = point_weights[:, :, offset:offset + chunk.shape[2], None, None] |
| occupancy = torch.exp(-distance / (2.0 * sigma * sigma)) * weights |
| weighted_maximum = torch.maximum(weighted_maximum, occupancy.amax(dim=2)) |
| offset += chunk.shape[2] |
| if point_weights is not None: |
| return weighted_maximum |
| return torch.exp(-minimum / (2.0 * sigma * sigma)) |
|
|
|
|
| def soft_rasterize_virtual_segments06( |
| coordinates: Tensor, state_logits: Tensor, *, size: int = 32, sigma: float = 0.025, |
| segment_stride: int = 2, |
| ) -> Tensor: |
| """필요 변수: 좌표·pen state·출력 크기. 작동 원리: pen-start 연결을 억제한 선분 거리로 differentiable raster를 만든다.""" |
|
|
| if coordinates.shape[:-1] != state_logits.shape[:-1] or state_logits.shape[-1] != 3: |
| raise ValueError("coordinate와 state logit의 batch·가설·시간 축이 일치해야 합니다.") |
| if size <= 0 or sigma <= 0 or segment_stride <= 0: |
| raise ValueError("raster size·sigma·segment_stride는 양수여야 합니다.") |
| axis = (torch.arange(size, device=coordinates.device, dtype=coordinates.dtype) + 0.5) / size |
| grid_y, grid_x = torch.meshgrid(axis, axis, indexing="ij") |
| grid = torch.stack((grid_x, grid_y), dim=-1) |
| starts = coordinates[:, :, :-segment_stride:segment_stride] |
| ends = coordinates[:, :, segment_stride::segment_stride] |
| segment_count = min(starts.shape[2], ends.shape[2]) |
| starts, ends = starts[:, :, :segment_count], ends[:, :, :segment_count] |
| state_probability = state_logits.softmax(dim=-1) |
| pen_start = state_probability[..., 1] |
| pen_end = state_probability[..., 2] |
| valid_rows = [] |
| for start in range(0, coordinates.shape[2] - segment_stride, segment_stride): |
| boundary = pen_start[..., start + 1:start + segment_stride + 1].amax(dim=-1) |
| |
| ended_before_target = pen_end[..., start:start + segment_stride + 1].amax(dim=-1) |
| valid_rows.append((1.0 - boundary) * (1.0 - ended_before_target)) |
| segment_valid = torch.stack(valid_rows[:segment_count], dim=2) |
| maximum = torch.zeros( |
| (*coordinates.shape[:2], size, size), dtype=coordinates.dtype, device=coordinates.device, |
| ) |
| for first in range(0, segment_count, 16): |
| start = starts[:, :, first:first + 16, None, None] |
| vector = (ends[:, :, first:first + 16] - starts[:, :, first:first + 16])[:, :, :, None, None] |
| relative = grid - start |
| projection = (relative * vector).sum(dim=-1) / vector.square().sum(dim=-1).clamp_min(1e-8) |
| closest = start + projection.clamp(0.0, 1.0).unsqueeze(-1) * vector |
| distance = (grid - closest).square().sum(dim=-1) |
| occupancy = torch.exp(-distance / (2.0 * sigma * sigma)) |
| occupancy = occupancy * segment_valid[:, :, first:first + 16, None, None] |
| maximum = torch.maximum(maximum, occupancy.amax(dim=2)) |
| |
| points = soft_rasterize_virtual06( |
| coordinates, size=size, sigma=sigma, point_stride=max(1, coordinates.shape[2] // 32), |
| point_weights=1.0 - pen_end, |
| ) |
| return torch.maximum(maximum, points) |
|
|
|
|
| def virtual_raster_similarity06( |
| coordinates: Tensor, raster: Tensor, *, state_logits: Tensor | None = None, |
| size: int = 32, sigma: float = 0.025, |
| ) -> Tensor: |
| """필요 변수: 가설 좌표·원본 raster. 작동 원리: 재렌더링 Dice와 양방향 coverage로 라벨 독립 품질을 계산한다.""" |
|
|
| reconstructed = ( |
| soft_rasterize_virtual_segments06(coordinates, state_logits, size=size, sigma=sigma) |
| if state_logits is not None else soft_rasterize_virtual06( |
| coordinates, size=size, sigma=sigma, point_stride=max(1, coordinates.shape[2] // 64), |
| ) |
| ) |
| target = nn.functional.adaptive_max_pool2d(raster, (size, size))[:, 0] |
| target = target[:, None].expand_as(reconstructed) |
| intersection = (reconstructed * target).sum(dim=(-1, -2)) |
| dice = (2.0 * intersection + 1e-5) / ( |
| reconstructed.sum(dim=(-1, -2)) + target.sum(dim=(-1, -2)) + 1e-5 |
| ) |
| |
| precision = intersection / reconstructed.sum(dim=(-1, -2)).clamp_min(1e-5) |
| recall = intersection / target.sum(dim=(-1, -2)).clamp_min(1e-5) |
| coverage = torch.sqrt((precision * recall).clamp_min(0.0)) |
| return 0.5 * (dice + coverage) |
|
|
|
|
| def raster_symmetry_logits06( |
| model: "MathInk06Model", output: dict[str, Tensor], *, mode: str = "logsumexp", |
| ) -> tuple[Tensor, Tensor]: |
| """필요 변수: 가상 stroke 출력·shared 모델. 작동 원리: 정적 이미지에서 알 수 없는 방향·획순서 네 경우를 동일 encoder로 평가한다.""" |
|
|
| features = virtual_features06( |
| output["coordinates"], output["state_logits"], output["stroke_progress"], |
| contract=model.virtual_contract, |
| ) |
| batch, hypotheses, steps, channels = features.shape |
| if model.use_virtual_adapter: |
| features = model.virtual_adapter(features.view(batch * hypotheses, steps, channels)).view( |
| batch, hypotheses, steps, channels, |
| ) |
| variants = equivalent_modality_features06(features.view(batch * hypotheses, steps, channels)) |
| exact, family = model.classify_trajectory(variants.flatten(0, 1)) |
| exact = exact.view(batch, hypotheses, 4, -1) |
| family = family.view(batch, hypotheses, 4, -1) |
| if mode == "logsumexp": |
| return torch.logsumexp(exact.log_softmax(dim=-1), dim=2), torch.logsumexp( |
| family.log_softmax(dim=-1), dim=2, |
| ) |
| if mode == "max": |
| return exact.log_softmax(dim=-1).amax(dim=2), family.log_softmax(dim=-1).amax(dim=2) |
| raise ValueError(f"지원하지 않는 symmetry mode입니다: {mode}") |
|
|
|
|
| class MathInk06Model(nn.Module): |
| """필요 변수: 378 exact/family class와 선택 boundary head. 작동 원리: online·virtual stroke를 동일 embedding으로 분류한다.""" |
|
|
| def __init__( |
| self, *, exact_classes: int, family_classes: int, hidden_size: int = 128, hypotheses: int = 4, |
| raster_architecture: str = "spatial_flat_progress_v2", virtual_contract: str = "legacy_v1", |
| use_virtual_adapter: bool = False, use_boundary_head: bool = False, |
| ) -> None: |
| super().__init__() |
| if raster_architecture not in { |
| "spatial_flat_v1", "spatial_flat_progress_v2", "cross_attention_v2", "cross_attention_8x8_v3", |
| "gated_cross_attention_8x8_v4", "split_auxiliary_v5", "fine_cross_attention_16x16_v6", |
| "gated_fine_cross_attention_16x16_v7", |
| "ink_pointer_32x32_v8", |
| }: |
| raise ValueError("지원하지 않는 raster architecture입니다.") |
| self.hidden_size = hidden_size |
| self.hypotheses = hypotheses |
| self.raster_architecture = raster_architecture |
| self.virtual_contract = virtual_contract |
| self.use_virtual_adapter = use_virtual_adapter |
| self.use_boundary_head = use_boundary_head |
| self.virtual_adapter_weight = 1.0 |
| self.trajectory_encoder = SharedTrajectoryEncoder06(hidden_size=hidden_size) |
| self.virtual_adapter = VirtualTrajectoryAdapter06() |
| self.exact_head = nn.Linear(hidden_size * 3, exact_classes) |
| self.family_head = nn.Linear(hidden_size * 3, family_classes) |
| self.boundary_head = nn.Linear(hidden_size * 3, 1) if use_boundary_head else None |
| self.raster_encoder = DepthwiseRasterEncoder06(hidden_size) |
| self.virtual_decoder = VirtualStrokeDecoder06(hidden_size, hypotheses) |
| self.auxiliary_virtual_decoder = ( |
| VirtualStrokeDecoder06(hidden_size, 2) if raster_architecture == "split_auxiliary_v5" else None |
| ) |
|
|
| def initialize_auxiliary_from_primary(self) -> None: |
| """필요 변수: split auxiliary 모델. 작동 원리: 유효한 primary 0·1번 출력을 auxiliary 초기값으로 복제한다.""" |
|
|
| if self.auxiliary_virtual_decoder is None: |
| raise ValueError("split_auxiliary_v5 모델에서만 auxiliary 초기화가 가능합니다.") |
| source = self.virtual_decoder.state_dict() |
| target = self.auxiliary_virtual_decoder.state_dict() |
| for key, target_value in target.items(): |
| source_value = source[key] |
| if source_value.shape == target_value.shape: |
| target[key] = source_value.detach().clone() |
| elif key == "hypothesis.weight" and source_value.shape[0] >= 2: |
| target[key] = source_value[:2].detach().clone() |
| else: |
| raise ValueError(f"auxiliary 초기화 shape가 일치하지 않습니다: {key}") |
| self.auxiliary_virtual_decoder.load_state_dict(target) |
|
|
| def encode_trajectory(self, sequence: Tensor) -> Tensor: |
| """필요 변수: B×128×19 canonical sequence. 작동 원리: 모든 symbol/behavior head가 공유할 trajectory embedding을 한 번 계산한다.""" |
|
|
| return self.trajectory_encoder(sequence) |
|
|
| def classify_trajectory(self, sequence: Tensor) -> tuple[Tensor, Tensor]: |
| """필요 변수: B×128×19. 작동 원리: shared embedding에서 기존 exact/family 출력 계약을 유지한다.""" |
|
|
| embedding = self.encode_trajectory(sequence) |
| return self.exact_head(embedding), self.family_head(embedding) |
|
|
| def classify_trajectory_with_boundary(self, sequence: Tensor) -> tuple[Tensor, Tensor, Tensor]: |
| """필요 변수: boundary head가 활성화된 sequence. 작동 원리: 한 embedding에서 exact/family/경계 침범 logit을 함께 반환한다.""" |
|
|
| if self.boundary_head is None: |
| raise RuntimeError("boundary head가 활성화되지 않았습니다.") |
| embedding = self.encode_trajectory(sequence) |
| return self.exact_head(embedding), self.family_head(embedding), self.boundary_head(embedding).squeeze(-1) |
|
|
| def forward_online(self, sequence: Tensor) -> tuple[Tensor, Tensor]: |
| """필요 변수: 실제 canonical tap. 작동 원리: raster 우회 없이 shared trajectory 분류를 반환한다.""" |
|
|
| return self.classify_trajectory(sequence) |
|
|
| def forward_online_with_boundary(self, sequence: Tensor) -> tuple[Tensor, Tensor, Tensor]: |
| """필요 변수: 실제 canonical tap. 작동 원리: 기존 LiteRT forward를 바꾸지 않고 연구용 boundary logit을 추가 노출한다.""" |
|
|
| return self.classify_trajectory_with_boundary(sequence) |
|
|
| def decode_raster_trajectories(self, raster: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]: |
| """필요 변수: B×1×128×128. 작동 원리: architecture별 top-4 좌표·state·progress·score를 한 경로로 만든다.""" |
|
|
| embedding, spatial_tokens = self.raster_encoder( |
| raster, fine_tokens=self.raster_architecture in { |
| "fine_cross_attention_16x16_v6", "gated_fine_cross_attention_16x16_v7", |
| }, |
| pointer_tokens=self.raster_architecture == "ink_pointer_32x32_v8", |
| ) |
| if self.raster_architecture == "split_auxiliary_v5": |
| if self.auxiliary_virtual_decoder is None: |
| raise RuntimeError("split auxiliary decoder가 초기화되지 않았습니다.") |
| primary = self.virtual_decoder(embedding) |
| auxiliary = self.auxiliary_virtual_decoder(embedding, spatial_tokens, gated_attention=True) |
| return tuple( |
| torch.cat((primary[index][:, :2], auxiliary[index]), dim=1) for index in range(4) |
| ) |
| memory = spatial_tokens if self.raster_architecture in { |
| "cross_attention_v2", "cross_attention_8x8_v3", "gated_cross_attention_8x8_v4", |
| "fine_cross_attention_16x16_v6", "gated_fine_cross_attention_16x16_v7", |
| "ink_pointer_32x32_v8", |
| } else None |
| pointer_mode = self.raster_architecture == "ink_pointer_32x32_v8" |
| return self.virtual_decoder( |
| embedding, memory, gated_attention=self.raster_architecture in { |
| "gated_cross_attention_8x8_v4", "gated_fine_cross_attention_16x16_v7", |
| }, |
| pointer_positions=self.raster_encoder.pointer_positions if pointer_mode else None, |
| ink_prior=( |
| nn.functional.adaptive_max_pool2d(raster, (32, 32)).flatten(2)[:, 0] |
| if pointer_mode else None |
| ), |
| ) |
|
|
| def forward_raster(self, raster: Tensor) -> dict[str, Tensor]: |
| """필요 변수: B×1×128×128. 작동 원리: top-4 가상 stroke를 만든 뒤 shared TCN으로만 분류한다.""" |
|
|
| coordinates, states, progress, hypothesis_scores = self.decode_raster_trajectories(raster) |
| features = virtual_features06( |
| coordinates, states, None if self.raster_architecture == "spatial_flat_v1" else progress, |
| contract=self.virtual_contract, |
| ) |
| batch, hypotheses, steps, channels = features.shape |
| if self.use_virtual_adapter: |
| raw_features = features |
| adapted_features = self.virtual_adapter(features.view(batch * hypotheses, steps, channels)).view( |
| batch, hypotheses, steps, channels, |
| ) |
| features = raw_features + self.virtual_adapter_weight * (adapted_features - raw_features) |
| flat_features = features.view(batch * hypotheses, steps, channels) |
| if self.boundary_head is None: |
| exact, family = self.classify_trajectory(flat_features) |
| boundary = None |
| else: |
| exact, family, boundary = self.classify_trajectory_with_boundary(flat_features) |
| output = { |
| "coordinates": coordinates, "state_logits": states, "stroke_progress": progress, |
| "hypothesis_scores": hypothesis_scores, |
| "exact_logits": exact.view(batch, hypotheses, -1), "family_logits": family.view(batch, hypotheses, -1), |
| } |
| if boundary is not None: |
| output["boundary_logits"] = boundary.view(batch, hypotheses) |
| return output |
|
|
| def forward(self, sequence: Tensor) -> tuple[Tensor, Tensor]: |
| """필요 변수: LiteRT용 online tensor. 작동 원리: 기본 forward를 online 경로로 고정한다.""" |
|
|
| return self.forward_online(sequence) |
|
|
|
|
| def boundary_auxiliary_loss06( |
| boundary_logits: Tensor, |
| boundary_targets: Tensor, |
| *, |
| positive_weight: float = 1.0, |
| sample_weight: Tensor | None = None, |
| ) -> Tensor: |
| """필요 변수: 후보별 경계 logit·0/1 target·선택 weight. 작동 원리: class imbalance를 보정한 binary auxiliary loss를 계산한다.""" |
|
|
| if boundary_logits.shape != boundary_targets.shape: |
| raise ValueError("boundary logit과 target shape가 다릅니다.") |
| if positive_weight <= 0.0: |
| raise ValueError("boundary positive weight는 0보다 커야 합니다.") |
| targets = boundary_targets.to(dtype=boundary_logits.dtype) |
| loss = nn.functional.binary_cross_entropy_with_logits( |
| boundary_logits, |
| targets, |
| pos_weight=torch.as_tensor(positive_weight, dtype=boundary_logits.dtype, device=boundary_logits.device), |
| reduction="none", |
| ) |
| if sample_weight is not None: |
| if sample_weight.shape != loss.shape: |
| raise ValueError("boundary sample weight shape가 다릅니다.") |
| normalized = sample_weight.to(loss).clamp_min(0.0) |
| return (loss * normalized).sum() / normalized.sum().clamp_min(1e-8) |
| return loss.mean() |
|
|
|
|
| def fuse_raster_logits06( |
| output: dict[str, Tensor], *, mode: str = "max", score_weight: float = 1.0, |
| family_weight: float = 0.0, geometry_weight: float = 0.0, |
| exact_family_index: Tensor | None = None, |
| ) -> tuple[Tensor, Tensor]: |
| """필요 변수: top-4 exact/family/quality logit. 작동 원리: 기호별 증거를 합치고 debug 대표 가설을 반환한다.""" |
|
|
| exact = output["exact_logits"].log_softmax(dim=-1) |
| score = output["hypothesis_scores"].log_softmax(dim=-1).unsqueeze(-1) |
| joint = exact + score_weight * score |
| if geometry_weight: |
| if "geometry_scores" not in output: |
| raise ValueError("geometry_weight를 사용할 때 geometry_scores가 필요합니다.") |
| geometry = output["geometry_scores"].clamp_min(1e-6).log().unsqueeze(-1) |
| joint = joint + geometry_weight * geometry |
| if family_weight: |
| if exact_family_index is None: |
| raise ValueError("family_weight를 사용할 때 exact_family_index가 필요합니다.") |
| family = output["family_logits"].log_softmax(dim=-1)[..., exact_family_index] |
| joint = joint + family_weight * family |
| if mode == "max": |
| fused = joint.amax(dim=1) |
| elif mode == "logsumexp": |
| fused = torch.logsumexp(joint, dim=1) |
| elif mode == "score_pick": |
| selected = output["hypothesis_scores"].argmax(dim=1) |
| fused = joint[torch.arange(len(joint), device=joint.device), selected] |
| return fused, selected |
| else: |
| raise ValueError(f"지원하지 않는 raster fusion mode입니다: {mode}") |
| |
| predicted = fused.argmax(dim=-1) |
| contribution = joint.gather(2, predicted[:, None, None].expand(-1, joint.shape[1], 1)).squeeze(-1) |
| return fused, contribution.argmax(dim=1) |
|
|
|
|
| def fuse_hypothesis_class_logits06( |
| class_logits: Tensor, hypothesis_scores: Tensor, *, mode: str = "logsumexp", score_weight: float = 1.0, |
| ) -> Tensor: |
| """필요 변수: B×H×C 분류 logit·B×H 가설 점수. 작동 원리: exact/family 공통 규칙으로 top-H 증거를 결합한다.""" |
|
|
| if class_logits.ndim != 3 or hypothesis_scores.shape != class_logits.shape[:2]: |
| raise ValueError("class logit과 hypothesis score shape가 일치하지 않습니다.") |
| joint = class_logits.log_softmax(dim=-1) |
| joint = joint + score_weight * hypothesis_scores.log_softmax(dim=-1).unsqueeze(-1) |
| if mode == "max": |
| return joint.amax(dim=1) |
| if mode == "logsumexp": |
| return torch.logsumexp(joint, dim=1) |
| if mode == "score_pick": |
| selected = hypothesis_scores.argmax(dim=1) |
| return joint[torch.arange(len(joint), device=joint.device), selected] |
| raise ValueError(f"지원하지 않는 hypothesis fusion mode입니다: {mode}") |
|
|
|
|
| def hypothesis_quality_features06(output: dict[str, Tensor], exact_family_index: Tensor) -> Tensor: |
| """필요 변수: 가설별 logits/state/좌표·family 사상. 작동 원리: raster label 없이 가설 품질 특징을 만든다.""" |
|
|
| exact_log_probability = output["exact_logits"].log_softmax(dim=-1) |
| exact_probability = exact_log_probability.exp() |
| top_values, top_indices = exact_log_probability.topk(min(2, exact_log_probability.shape[-1]), dim=-1) |
| predicted = top_indices[..., 0] |
| margin = top_values[..., 0] - top_values[..., -1] |
| exact_entropy = -(exact_probability * exact_log_probability).sum(dim=-1) / np.log(max(2, exact_probability.shape[-1])) |
| family_log_probability = output["family_logits"].log_softmax(dim=-1) |
| predicted_family = exact_family_index[predicted] |
| predicted_family_log_probability = family_log_probability.gather(2, predicted_family.unsqueeze(-1)).squeeze(-1) |
| state_log_probability = output["state_logits"].log_softmax(dim=-1) |
| state_probability = state_log_probability.exp() |
| state_entropy = -(state_probability * state_log_probability).sum(dim=-1).mean(dim=-1) / np.log(3.0) |
| start_confidence = state_probability[..., 1].amax(dim=-1) |
| delta = output["coordinates"][:, :, 1:] - output["coordinates"][:, :, :-1] |
| distance = delta.square().sum(dim=-1).sqrt() |
| path_length = distance.mean(dim=-1) |
| direction = delta / distance.clamp_min(1e-6).unsqueeze(-1) |
| turn = ( |
| direction[:, :, :-1, 0] * direction[:, :, 1:, 1] |
| - direction[:, :, :-1, 1] * direction[:, :, 1:, 0] |
| ).abs().mean(dim=-1) |
| agreement = (predicted[:, :, None] == predicted[:, None, :]).float().mean(dim=-1) |
| branch = nn.functional.one_hot( |
| torch.arange(predicted.shape[1], device=predicted.device), num_classes=predicted.shape[1], |
| ).to(dtype=exact_probability.dtype).unsqueeze(0).expand(len(predicted), -1, -1) |
| scalar = torch.stack(( |
| top_values[..., 0], margin, exact_entropy, family_log_probability.amax(dim=-1), |
| predicted_family_log_probability, output["hypothesis_scores"].log_softmax(dim=-1), |
| start_confidence, state_entropy, path_length, turn, agreement, |
| ), dim=-1) |
| return torch.cat((scalar, branch), dim=-1) |
|
|
|
|
| class HypothesisSelector06(nn.Module): |
| """필요 변수: trajectory-only 품질 특징. 작동 원리: 각 virtual hypothesis의 혼합 logit을 예측한다.""" |
|
|
| def __init__( |
| self, input_size: int = 15, hidden_size: int = 24, *, label_classes: int = 0, |
| label_embedding_size: int = 0, |
| ) -> None: |
| super().__init__() |
| if (label_classes > 0) != (label_embedding_size > 0): |
| raise ValueError("label class와 embedding 크기는 함께 지정해야 합니다.") |
| self.label_embedding = ( |
| nn.Embedding(label_classes, label_embedding_size) if label_classes > 0 else None |
| ) |
| self.network = nn.Sequential( |
| nn.LayerNorm(input_size + label_embedding_size), |
| nn.Linear(input_size + label_embedding_size, hidden_size), nn.GELU(), |
| nn.Linear(hidden_size, 1), |
| ) |
|
|
| def forward(self, features: Tensor, predicted_labels: Tensor | None = None) -> Tensor: |
| """필요 변수: 품질 특징·선택 top-1 label. 작동 원리: 가설별 scalar quality logit을 반환한다.""" |
|
|
| if self.label_embedding is not None: |
| if predicted_labels is None: |
| raise ValueError("class-conditional selector에는 predicted_labels가 필요합니다.") |
| features = torch.cat((features, self.label_embedding(predicted_labels)), dim=-1) |
| return self.network(features).squeeze(-1) |
|
|
|
|
| def initialize_from_05(model: MathInk06Model, checkpoint_paths: Sequence[Path]) -> None: |
| """필요 변수: 0.6 모델·동일 0.5 seed checkpoint. 작동 원리: 유효한 단일 teacher를 19채널 student 초기값으로 이식한다.""" |
|
|
| if not checkpoint_paths: |
| raise ValueError("0.5 checkpoint가 필요합니다.") |
| |
| |
| checkpoint = torch.load(checkpoint_paths[0], map_location="cpu", weights_only=False) |
| state = checkpoint["state_dict"] |
| target = model.state_dict() |
| mapping = { |
| "trajectory_encoder.input_projection": "encoder.input_projection", |
| "trajectory_encoder.blocks": "encoder.blocks", |
| "trajectory_encoder.attention": "encoder.attention", |
| "exact_head": "exact_head", "family_head": "family_head", |
| } |
| for target_key in list(target): |
| source_key = next((target_key.replace(prefix, source) for prefix, source in mapping.items() if target_key.startswith(prefix)), None) |
| if source_key is None or source_key not in state: |
| continue |
| source_value = state[source_key].float() |
| if source_value.shape == target[target_key].shape: |
| target[target_key] = source_value |
| elif target_key.endswith("input_projection.0.weight") and source_value.shape[1] == 15 and target[target_key].shape[1] == 19: |
| expanded = torch.zeros_like(target[target_key]) |
| expanded[:, :15] = source_value |
| target[target_key] = expanded |
| model.load_state_dict(target) |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class SymbolCandidate06: |
| """필요 변수: token·확률. 작동 원리: 모바일 공개 결과의 후보 한 개를 표현한다.""" |
|
|
| token: str |
| probability: float |
|
|
|
|
| class MathInk06Engine: |
| """필요 변수: 0.6 checkpoint. 작동 원리: 원본 stroke 또는 raster에서 텍스트 후보만 반환한다.""" |
|
|
| def __init__( |
| self, checkpoint: Path, *, adapter_checkpoint: Path | None = None, device: str = "cpu", |
| ) -> None: |
| """필요 변수: base와 선택 composite adapter. 작동 원리: base→shared state→modality adapter 순서로 런타임을 구성한다.""" |
|
|
| payload = torch.load(checkpoint, map_location=device, weights_only=False) |
| self.labels = tuple(str(value) for value in payload["exact_labels"]) |
| self.family_labels = tuple(str(value) for value in payload["family_labels"]) |
| self.model_version = str(payload.get("model_version", "aiflow-math-ink-0.6")) |
| self.model = MathInk06Model( |
| exact_classes=len(self.labels), family_classes=len(payload["family_labels"]), |
| hidden_size=int(payload["hidden_size"]), hypotheses=int(payload.get("hypotheses", 4)), |
| raster_architecture=str(payload.get("raster_architecture", "spatial_flat_v1")), |
| virtual_contract=str(payload.get("virtual_contract", "legacy_v1")), |
| use_virtual_adapter=bool(payload.get("use_virtual_adapter", False)), |
| ).to(device) |
| |
| self.model.load_state_dict(payload["state_dict"], strict=False) |
| self.model.virtual_adapter_weight = float(payload.get("virtual_adapter_weight", 1.0)) |
| self.model.eval() |
| self.device = torch.device(device) |
| self.raster_fusion = { |
| "mode": "max", "score_weight": 1.0, "family_weight": 0.0, "geometry_weight": 0.0, |
| "symmetry_weight": 0.0, "symmetry_mode": "logsumexp", |
| **dict(payload.get("raster_fusion", {})), |
| } |
| self.composite_adapter: nn.Module = nn.Identity() |
| self.online_adapter: nn.Module = nn.Identity() |
| self.raster_adapter: nn.Module = nn.Identity() |
| self.online_family_fusion_weight = 0.0 |
| if adapter_checkpoint is not None: |
| adapter_payload = torch.load( |
| adapter_checkpoint, map_location=device, weights_only=False, |
| ) |
| shared_state = adapter_payload.get("shared_state_dict") or {} |
| if shared_state: |
| incompatible = self.model.load_state_dict(shared_state, strict=False) |
| if incompatible.unexpected_keys: |
| raise ValueError( |
| f"adapter shared state key 오류: {incompatible.unexpected_keys}" |
| ) |
| architecture = str(adapter_payload["adapter_architecture"]) |
| if architecture == "local_v1": |
| adapter: nn.Module = VirtualTrajectoryAdapter06() |
| else: |
| from .skeleton_adapter06 import ( |
| DualModalityTrajectoryAdapter06, SkeletonTrajectoryAdapter06, |
| ) |
| if architecture == "tcn_v2": |
| adapter = SkeletonTrajectoryAdapter06() |
| elif architecture == "dual_tcn_v3": |
| adapter = DualModalityTrajectoryAdapter06() |
| else: |
| raise ValueError( |
| f"지원하지 않는 adapter architecture입니다: {architecture}" |
| ) |
| adapter.load_state_dict(adapter_payload["state_dict"]) |
| adapter = adapter.to(self.device).eval() |
| self.online_family_fusion_weight = float( |
| adapter_payload.get("family_fusion_weight", 0.0) |
| ) |
| if not 0.0 <= self.online_family_fusion_weight <= 1.0: |
| raise ValueError("online family fusion weight는 0~1 범위여야 합니다.") |
| self.composite_adapter = adapter |
| if architecture == "dual_tcn_v3": |
| self.online_adapter = adapter.online |
| self.raster_adapter = adapter.raster |
| else: |
| self.online_adapter = self.raster_adapter = adapter |
| self.model_version = ( |
| f"{self.model_version}+{adapter_payload.get('model_version', architecture)}" |
| ) |
| family_to_index = {label: index for index, label in enumerate(self.family_labels)} |
| |
| from .trajectory_sequence import shape_family |
| self.exact_family_index = torch.tensor( |
| [family_to_index[shape_family(label)] for label in self.labels], device=self.device, |
| ) |
| selector_payload = payload.get("hypothesis_selector") |
| self.hypothesis_selector: HypothesisSelector06 | None = None |
| if selector_payload: |
| self.hypothesis_selector = HypothesisSelector06( |
| input_size=int(selector_payload["input_size"]), hidden_size=int(selector_payload["hidden_size"]), |
| label_classes=int(selector_payload.get("label_classes", 0)), |
| label_embedding_size=int(selector_payload.get("label_embedding_size", 0)), |
| ).to(self.device) |
| self.hypothesis_selector.load_state_dict(selector_payload["state_dict"]) |
| self.hypothesis_selector.eval() |
|
|
| def fuse_raster_output(self, output: dict[str, Tensor]) -> tuple[Tensor, Tensor]: |
| """필요 변수: model raster 출력. 작동 원리: checkpoint에 따라 learned selector 또는 고정 fusion을 적용한다.""" |
|
|
| if self.hypothesis_selector is None: |
| return fuse_raster_logits06( |
| output, mode=str(self.raster_fusion["mode"]), |
| score_weight=float(self.raster_fusion["score_weight"]), |
| family_weight=float(self.raster_fusion["family_weight"]), |
| geometry_weight=float(self.raster_fusion["geometry_weight"]), |
| exact_family_index=self.exact_family_index, |
| ) |
| features = hypothesis_quality_features06(output, self.exact_family_index) |
| predicted_labels = output["exact_logits"].argmax(dim=-1) |
| selector_log_probability = self.hypothesis_selector(features, predicted_labels).log_softmax(dim=1) |
| joint = output["exact_logits"].log_softmax(dim=-1) + selector_log_probability.unsqueeze(-1) |
| fused = torch.logsumexp(joint, dim=1) |
| predicted = fused.argmax(dim=-1) |
| contribution = joint.gather(2, predicted[:, None, None].expand(-1, joint.shape[1], 1)).squeeze(-1) |
| return fused, contribution.argmax(dim=1) |
|
|
| def _forward_raster_composite06(self, raster: Tensor) -> dict[str, Tensor]: |
| """필요 변수: 정규화 raster. 작동 원리: virtual top-4를 외부 raster adapter까지 거쳐 shared head로 분류한다.""" |
|
|
| if isinstance(self.raster_adapter, nn.Identity): |
| return self.model.forward_raster(raster) |
| 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 = self.raster_adapter( |
| features.reshape(batch * hypotheses, steps, channels), |
| ) |
| exact, family = self.model.classify_trajectory(flat) |
| return { |
| "coordinates": coordinates, |
| "state_logits": states, |
| "stroke_progress": progress, |
| "hypothesis_scores": hypothesis_scores, |
| "exact_logits": exact.reshape(batch, hypotheses, -1), |
| "family_logits": family.reshape(batch, hypotheses, -1), |
| } |
|
|
| def _result(self, logits: Tensor, started: float, top_k: int) -> dict[str, Any]: |
| """필요 변수: fused logits·시작시각·k. 작동 원리: stroke/image 없이 모바일 공개 SymbolResult를 만든다.""" |
|
|
| probability = logits.softmax(dim=-1)[0] |
| values, indices = probability.topk(min(top_k, len(self.labels))) |
| candidates = [SymbolCandidate06(self.labels[int(index)], float(value)) for value, index in zip(values, indices, strict=True)] |
| return { |
| "candidates": [candidate.__dict__ if hasattr(candidate, "__dict__") else {"token": candidate.token, "probability": candidate.probability} for candidate in candidates], |
| "confidence": candidates[0].probability, "modelVersion": self.model_version, |
| "latencyMs": (time.perf_counter() - started) * 1000.0, |
| } |
|
|
| def _fuse_online_exact06(self, exact: Tensor, family: Tensor) -> Tensor: |
| """필요 변수: exact/family logit. 작동 원리: validation에서 고정한 형태군 prior로 exact 후보만 재정렬한다.""" |
|
|
| if not self.online_family_fusion_weight: |
| return exact |
| return ( |
| exact.log_softmax(dim=-1) |
| + self.online_family_fusion_weight |
| * family.log_softmax(dim=-1)[:, self.exact_family_index] |
| ) |
|
|
| def recognize_online(self, strokes: Sequence[dict[str, Any]], *, canvas_width: float, canvas_height: float, top_k: int = 5) -> dict[str, Any]: |
| """필요 변수: 원본 stroke·canvas. 작동 원리: 6Hz 재구성 후 기기 밖으로 내보낼 텍스트 후보만 반환한다.""" |
|
|
| started = time.perf_counter() |
| ink = canonicalize_ink06(strokes, canvas_width=canvas_width, canvas_height=canvas_height) |
| sequence = torch.from_numpy(ink.features).unsqueeze(0).to(self.device) |
| with torch.inference_mode(): |
| exact, family = self.model.forward_online(self.online_adapter(sequence)) |
| exact = self._fuse_online_exact06(exact, family) |
| return self._result(exact, started, top_k) |
|
|
| def recognize_raster(self, image: Image.Image, *, top_k: int = 5, debug: bool = False) -> dict[str, Any]: |
| """필요 변수: PIL image·k·로컬 debug. 작동 원리: 가상 stroke를 거쳐 텍스트만 반환하고 debug 때만 좌표를 붙인다.""" |
|
|
| started = time.perf_counter() |
| normalized = image.convert("L").resize((128, 128), Image.Resampling.LANCZOS) |
| raster = 1.0 - torch.from_numpy(np.asarray(normalized, dtype=np.float32) / 255.0) |
| with torch.inference_mode(): |
| output = self._forward_raster_composite06( |
| raster.view(1, 1, 128, 128).to(self.device), |
| ) |
| symmetry_weight = float(self.raster_fusion["symmetry_weight"]) |
| if symmetry_weight: |
| symmetry_exact, symmetry_family = raster_symmetry_logits06( |
| self.model, output, mode=str(self.raster_fusion["symmetry_mode"]), |
| ) |
| output["exact_logits"] = ( |
| (1.0 - symmetry_weight) * output["exact_logits"].log_softmax(dim=-1) |
| + symmetry_weight * symmetry_exact |
| ) |
| output["family_logits"] = ( |
| (1.0 - symmetry_weight) * output["family_logits"].log_softmax(dim=-1) |
| + symmetry_weight * symmetry_family |
| ) |
| if float(self.raster_fusion["geometry_weight"]): |
| output["geometry_scores"] = virtual_raster_similarity06( |
| output["coordinates"], raster.view(1, 1, 128, 128).to(self.device), |
| ) |
| logits, selected = self.fuse_raster_output(output) |
| flat_index = int(selected[0]) |
| result = self._result(logits, started, top_k) |
| if debug: |
| result["virtualHypothesis"] = output["coordinates"][0, flat_index].cpu().tolist() |
| return result |
|
|
