Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from typing import Any | |
| import torch | |
| from torch import nn | |
| from transformers import AutoModel | |
| from tiny_router.config import RouterModelConfig | |
| from tiny_router.constants import HEAD_LABELS | |
| class TinyRouterModel(nn.Module): | |
| def __init__(self, config: RouterModelConfig): | |
| super().__init__() | |
| self.config = config | |
| self.encoder = AutoModel.from_pretrained(config.encoder_name) | |
| hidden_size = self.encoder.config.hidden_size | |
| if config.pooling_type not in {"mean", "attention"}: | |
| raise ValueError(f"Unsupported pooling type: {config.pooling_type}") | |
| self.dropout = nn.Dropout(config.dropout) | |
| self.pooling_type = config.pooling_type | |
| if self.pooling_type == "attention": | |
| self.attention_pool = nn.Linear(hidden_size, 1) | |
| self.action_embedding = nn.Embedding(len(config.action_vocab), config.structured_hidden_dim) | |
| self.outcome_embedding = nn.Embedding(len(config.outcome_vocab), config.structured_hidden_dim) | |
| self.recency_projection = nn.Sequential( | |
| nn.Linear(1, config.recency_embed_dim), | |
| nn.ReLU(), | |
| nn.Linear(config.recency_embed_dim, config.recency_embed_dim), | |
| ) | |
| structured_dim = (config.structured_hidden_dim * 2) + config.recency_embed_dim | |
| canonical_head_order = list(HEAD_LABELS.keys()) | |
| extra_heads = [head for head in config.label_maps if head not in HEAD_LABELS] | |
| self.head_order = [head for head in canonical_head_order if head in config.label_maps] + extra_heads | |
| self.head_dependencies = { | |
| head: self.head_order[:idx] | |
| for idx, head in enumerate(self.head_order) | |
| } | |
| self.dependency_projections = nn.ModuleDict() | |
| self.heads = nn.ModuleDict( | |
| self._build_heads(hidden_size + structured_dim) | |
| ) | |
| self.loss_fns = {head: nn.CrossEntropyLoss() for head in config.label_maps} | |
| def _build_heads(self, base_dim: int) -> dict[str, nn.Linear]: | |
| heads: dict[str, nn.Linear] = {} | |
| for head, labels in self.config.label_maps.items(): | |
| input_dim = base_dim | |
| if self.config.use_head_dependencies and self.head_dependencies[head]: | |
| dependency_input_dim = sum( | |
| len(self.config.label_maps[source]) | |
| for source in self.head_dependencies[head] | |
| ) | |
| self.dependency_projections[head] = nn.Sequential( | |
| nn.Linear(dependency_input_dim, self.config.dependency_hidden_dim), | |
| nn.Tanh(), | |
| nn.Dropout(self.config.dropout), | |
| ) | |
| input_dim += self.config.dependency_hidden_dim | |
| heads[head] = nn.Linear(input_dim, len(labels)) | |
| return heads | |
| def pool_hidden_state(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: | |
| if self.pooling_type == "attention": | |
| scores = self.attention_pool(hidden_state).squeeze(-1) | |
| scores = scores.masked_fill(attention_mask == 0, -1e4) | |
| weights = torch.softmax(scores, dim=-1).unsqueeze(-1) | |
| return torch.sum(hidden_state * weights, dim=1) | |
| mask = attention_mask.unsqueeze(-1).expand(hidden_state.size()).float() | |
| summed = torch.sum(hidden_state * mask, dim=1) | |
| counts = torch.clamp(mask.sum(dim=1), min=1e-6) | |
| return summed / counts | |
| def forward( | |
| self, | |
| input_ids: torch.Tensor, | |
| attention_mask: torch.Tensor, | |
| previous_action_id: torch.Tensor, | |
| previous_outcome_id: torch.Tensor, | |
| log_recency_seconds: torch.Tensor, | |
| has_interaction: torch.Tensor, | |
| has_recency: torch.Tensor | None = None, | |
| head_loss_weights: dict[str, float] | None = None, | |
| **kwargs: Any, | |
| ) -> dict[str, Any]: | |
| outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask) | |
| pooled = self.pool_hidden_state(outputs.last_hidden_state, attention_mask) | |
| pooled = self.dropout(pooled) | |
| structured_mask = has_interaction.float().unsqueeze(-1) | |
| recency_mask_source = has_interaction if has_recency is None else has_recency | |
| recency_mask = recency_mask_source.float().unsqueeze(-1) | |
| action_emb = self.action_embedding(previous_action_id) * structured_mask | |
| outcome_emb = self.outcome_embedding(previous_outcome_id) * structured_mask | |
| recency_emb = self.recency_projection(log_recency_seconds.unsqueeze(-1)) * recency_mask | |
| combined = torch.cat([pooled, action_emb, outcome_emb, recency_emb], dim=-1) | |
| logits = {} | |
| prior_head_probs: dict[str, torch.Tensor] = {} | |
| for head in self.head_order: | |
| head_inputs = [combined] | |
| if self.config.use_head_dependencies and self.head_dependencies[head]: | |
| dependency_input = torch.cat( | |
| [prior_head_probs[source] for source in self.head_dependencies[head]], | |
| dim=-1, | |
| ) | |
| dependency_features = self.dependency_projections[head](dependency_input) | |
| head_inputs.append(dependency_features) | |
| classifier_input = torch.cat(head_inputs, dim=-1) | |
| head_logits = self.heads[head](classifier_input) | |
| logits[head] = head_logits | |
| # Keep the supervision for earlier heads local to their own loss while still | |
| # exposing their predictions to later heads. | |
| prior_head_probs[head] = torch.softmax(head_logits.detach(), dim=-1) | |
| result: dict[str, Any] = {"logits": logits} | |
| total_loss = None | |
| for head in self.config.label_maps: | |
| label_key = f"labels_{head}" | |
| if label_key not in kwargs: | |
| continue | |
| loss = self.loss_fns[head](logits[head], kwargs[label_key]) | |
| if head_loss_weights: | |
| loss = loss * float(head_loss_weights.get(head, 1.0)) | |
| total_loss = loss if total_loss is None else total_loss + loss | |
| result["loss"] = total_loss | |
| return result | |