OptimizerStudy/NCPL-intermediate

1import torch2import torch.nn as nn3from transformers import AutoModel, AutoConfig4 5 6class ScalingLawForecaster(nn.Module):7    def __init__(8        self,9        base_model_name: str = "HuggingFaceTB/SmolLM2-135M",10        init_from_pretrained: bool = True,11        force_fp32: bool = False,12    ):13        super().__init__()14        self.config = AutoConfig.from_pretrained(base_model_name)15        if force_fp32:16            self.config.torch_dtype = torch.float3217        if init_from_pretrained:18            if force_fp32:19                self.base = AutoModel.from_pretrained(20                    base_model_name,21                    config=self.config,22                    torch_dtype=torch.float32,23                )24            else:25                self.base = AutoModel.from_pretrained(base_model_name, config=self.config)26        else:27            self.base = AutoModel.from_config(self.config)28 29        hidden_size = self.config.hidden_size30 31        act_cls = nn.ReLU32        self.num_mlp = nn.Sequential(33            nn.Linear(1, hidden_size * 2),34            act_cls(),35            nn.Linear(hidden_size * 2, hidden_size)36        )37 38        self.head = nn.Linear(hidden_size, 1)39 40    def forward(41        self,42        input_ids: torch.LongTensor,43        is_number_mask: torch.BoolTensor,44        number_values_filled: torch.FloatTensor,45        attention_mask: torch.BoolTensor = None46    ) -> torch.FloatTensor:47        """48        Args:49            input_ids:          (batch, seq_len)50            is_number_mask:     (batch, seq_len)    bool mask for numeric tokens51            number_values_filled:(batch, seq_len)    float values (0 for non-numeric)52            attention_mask:     (batch, seq_len)    optional53        Returns:54            logits: (batch, seq_len) scalar predictions per token55        """56        # Text embeddings57        input_ids[input_ids == 49152] = 0 58        text_emb = self.base.get_input_embeddings()(input_ids)59 60        # Numeric MLP embeddings61        flat_vals = number_values_filled.view(-1, 1)62        mlp_out = self.num_mlp(flat_vals)  63        mlp_out = mlp_out.view_as(text_emb) 64 65        mask = is_number_mask.unsqueeze(-1)66        inputs_embeds = torch.where(mask, mlp_out, text_emb)67 68        outputs = self.base(69            inputs_embeds=inputs_embeds,70            attention_mask=attention_mask,71            return_dict=True72        )73        hidden = outputs.last_hidden_state 74 75        # Final scalar head76        logits = self.head(hidden).squeeze(-1)  77        return logits78 79