File size: 2,811 Bytes
d195287 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | import torch
import torch.nn as nn
class QuantOHLCEmbedder(nn.Module):
def __init__(
self,
num_features: int,
sequence_length: int = 60,
version_vocab_size: int = 4,
hidden_dim: int = 320,
num_layers: int = 3,
num_heads: int = 8,
output_dim: int = 1536,
dtype: torch.dtype = torch.float16,
):
super().__init__()
self.num_features = num_features
self.sequence_length = sequence_length
self.output_dim = output_dim
self.dtype = dtype
self.feature_proj = nn.Sequential(
nn.LayerNorm(num_features),
nn.Linear(num_features, hidden_dim),
nn.GELU(),
)
self.position_embedding = nn.Parameter(torch.zeros(1, sequence_length, hidden_dim))
self.version_embedding = nn.Embedding(version_vocab_size, hidden_dim, padding_idx=0)
encoder_layer = nn.TransformerEncoderLayer(
d_model=hidden_dim,
nhead=num_heads,
dim_feedforward=hidden_dim * 4,
dropout=0.0,
batch_first=True,
activation="gelu",
norm_first=True,
)
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.output_head = nn.Sequential(
nn.LayerNorm(hidden_dim),
nn.Linear(hidden_dim, hidden_dim * 2),
nn.GELU(),
nn.LayerNorm(hidden_dim * 2),
nn.Linear(hidden_dim * 2, output_dim),
nn.LayerNorm(output_dim),
)
self.to(dtype)
def forward(
self,
feature_tokens: torch.Tensor,
feature_mask: torch.Tensor,
version_ids: torch.Tensor,
) -> torch.Tensor:
if feature_tokens.ndim != 3:
raise ValueError(f"Expected [B, T, F], got {feature_tokens.shape}")
if feature_tokens.shape[1] != self.sequence_length:
raise ValueError(f"Expected T={self.sequence_length}, got {feature_tokens.shape[1]}")
if feature_tokens.shape[2] != self.num_features:
raise ValueError(f"Expected F={self.num_features}, got {feature_tokens.shape[2]}")
x = self.feature_proj(feature_tokens.to(self.dtype))
version_embed = self.version_embedding(version_ids).unsqueeze(1)
x = x + self.position_embedding[:, : x.shape[1], :].to(x.dtype) + version_embed
key_padding_mask = ~(feature_mask > 0)
x = self.encoder(x, src_key_padding_mask=key_padding_mask)
mask = feature_mask.to(x.dtype).unsqueeze(-1)
valid_any = (feature_mask.sum(dim=1, keepdim=True) > 0).to(x.dtype)
denom = mask.sum(dim=1).clamp_min(1.0)
pooled = (x * mask).sum(dim=1) / denom
out = self.output_head(pooled)
return out * valid_any
|