File size: 2,565 Bytes
49263e6 | 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 | import torch
import torch.nn as nn
from torch.nn import functional as F
from transformers import PretrainedConfig, PreTrainedModel
class TechcodeXConfig(PretrainedConfig):
model_type = "techcodex_hybrid_transformer"
def __init__(self, vocab_size=50257, n_embd=256, n_head=4, block_size=128, **kwargs):
super().__init__(**kwargs)
self.vocab_size = vocab_size
self.n_embd = n_embd
self.n_head = n_head
self.block_size = block_size
class ProprietaryRecurrentLayer(nn.Module):
def __init__(self, n_embd):
super().__init__()
self.hidden_dim = n_embd
self.gate_mix = nn.Linear(n_embd * 2, n_embd)
self.gate_state = nn.Linear(n_embd * 2, n_embd)
self.ln = nn.LayerNorm(n_embd)
def forward(self, x):
B, T, C = x.shape
hidden = torch.zeros(B, self.hidden_dim, device=x.device)
outputs = []
for t in range(T):
current_token = x[:, t, :]
combined = torch.cat([current_token, hidden], dim=-1)
mix = torch.sigmoid(self.gate_mix(combined))
new_state = torch.tanh(self.gate_state(combined))
hidden = (mix * hidden) + ((1.0 - mix) * new_state)
outputs.append(hidden.unsqueeze(1))
return self.ln(torch.cat(outputs, dim=1))
class TechcodeXModel(PreTrainedModel):
config_class = TechcodeXConfig
def __init__(self, config):
super().__init__(config)
self.config = config
self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
self.attn_layer = nn.TransformerEncoderLayer(
d_model=config.n_embd, nhead=config.n_head, dim_feedforward=config.n_embd*4, batch_first=True
)
self.recurrent_layer = ProprietaryRecurrentLayer(config.n_embd)
self.bridge = nn.Linear(config.n_embd * 2, config.n_embd)
self.ln_final = nn.LayerNorm(config.n_embd)
self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
self.post_init()
def forward(self, input_ids, labels=None, **kwargs):
B, T = input_ids.shape
positions = torch.arange(0, T, device=input_ids.device).unsqueeze(0)
x = self.token_embedding(input_ids) + self.position_embedding(positions)
path_a = self.attn_layer(x)
path_b = self.recurrent_layer(x)
x = self.bridge(torch.cat([path_a, path_b], dim=-1))
logits = self.lm_head(self.ln_final(x))
return {"logits": logits}
|