AgentA123/project-helios-api

1import math2 3import torch4from torch import nn5 6 7class PositionalEncoding(nn.Module):8    def __init__(self, d_model: int, dropout: float = 0.15, max_len: int = 5000):9        super().__init__()10        self.dropout = nn.Dropout(p=dropout)11 12        position = torch.arange(max_len).unsqueeze(1)13        div_term = torch.exp(14            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)15        )16 17        pe = torch.zeros(max_len, d_model)18        pe[:, 0::2] = torch.sin(position * div_term)19        pe[:, 1::2] = torch.cos(position * div_term)20        self.register_buffer("pe", pe.unsqueeze(0))21 22    def forward(self, x: torch.Tensor) -> torch.Tensor:23        x = x + self.pe[:, : x.size(1)]24        return self.dropout(x)25 26 27class SolarTransformer(nn.Module):28    def __init__(29        self,30        input_dim: int,31        d_model: int = 64,32        nhead: int = 4,33        num_layers: int = 3,34        dropout: float = 0.15,35        max_len: int = 3000,36        dim_feedforward: int = 128,37    ):38        super().__init__()39        self.input_proj = nn.Linear(input_dim, d_model)40        self.pos = PositionalEncoding(41            d_model=d_model,42            dropout=dropout,43            max_len=max_len,44        )45 46        encoder_layer = nn.TransformerEncoderLayer(47            d_model=d_model,48            nhead=nhead,49            dim_feedforward=dim_feedforward,50            dropout=dropout,51            batch_first=True,52        )53        self.encoder = nn.TransformerEncoder(54            encoder_layer=encoder_layer,55            num_layers=num_layers,56        )57        self.head = nn.Sequential(58            nn.LayerNorm(d_model),59            nn.Linear(d_model, 32),60            nn.ReLU(),61            nn.Dropout(dropout),62            nn.Linear(32, 1),63        )64 65    def forward(self, x: torch.Tensor) -> torch.Tensor:66        x = self.input_proj(x)67        x = self.pos(x)68        x = self.encoder(x)69        pooled = x.mean(dim=1)70        return self.head(pooled).squeeze(-1)71