|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
from transformers.modeling_outputs import BaseModelOutputWithPast |
|
|
from transformers import PreTrainedModel |
|
|
from configuration_theta import ThetaConfig |
|
|
|
|
|
class ThetaAttention(nn.Module): |
|
|
def __init__(self, config: ThetaConfig): |
|
|
super().__init__() |
|
|
self.num_heads = config.num_attention_heads |
|
|
self.head_dim = config.hidden_size // config.num_attention_heads |
|
|
self.scale = self.head_dim ** -0.5 |
|
|
|
|
|
self.q_proj = nn.Linear(config.hidden_size, config.hidden_size) |
|
|
self.k_proj = nn.Linear(config.hidden_size, config.hidden_size) |
|
|
self.v_proj = nn.Linear(config.hidden_size, config.hidden_size) |
|
|
self.out_proj = nn.Linear(config.hidden_size, config.hidden_size) |
|
|
|
|
|
def forward(self, hidden_states): |
|
|
batch_size, seq_length, embed_dim = hidden_states.size() |
|
|
|
|
|
query = self.q_proj(hidden_states) |
|
|
key = self.k_proj(hidden_states) |
|
|
value = self.v_proj(hidden_states) |
|
|
|
|
|
query = query.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) |
|
|
key = key.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) |
|
|
value = value.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) |
|
|
|
|
|
attn_scores = torch.matmul(query, key.transpose(-2, -1)) * self.scale |
|
|
attn_probs = nn.functional.softmax(attn_scores, dim=-1) |
|
|
attn_output = torch.matmul(attn_probs, value) |
|
|
|
|
|
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_length, embed_dim) |
|
|
attn_output = self.out_proj(attn_output) |
|
|
|
|
|
return attn_output |
|
|
|
|
|
class ThetaMLP(nn.Module): |
|
|
def __init__(self, config: ThetaConfig): |
|
|
super().__init__() |
|
|
self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) |
|
|
self.act = nn.SiLU() |
|
|
self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) |
|
|
|
|
|
def forward(self, hidden_states): |
|
|
hidden_states = self.fc1(hidden_states) |
|
|
hidden_states = self.act(hidden_states) |
|
|
hidden_states = self.fc2(hidden_states) |
|
|
return hidden_states |
|
|
|
|
|
class ThetaBlock(nn.Module): |
|
|
def __init__(self, config: ThetaConfig): |
|
|
super().__init__() |
|
|
self.attention = ThetaAttention(config) |
|
|
self.mlp = ThetaMLP(config) |
|
|
self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps) |
|
|
self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps) |
|
|
|
|
|
def forward(self, hidden_states): |
|
|
hidden_states = hidden_states + self.attention(self.norm1(hidden_states)) |
|
|
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) |
|
|
return hidden_states |
|
|
|
|
|
class ThetaModel(PreTrainedModel): |
|
|
config_class = ThetaConfig |
|
|
|
|
|
def __init__(self, config: ThetaConfig): |
|
|
super().__init__(config) |
|
|
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
|
|
self.layers = nn.ModuleList([ThetaBlock(config) for _ in range(config.num_hidden_layers)]) |
|
|
self.norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps) |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
input_ids=None, |
|
|
attention_mask=None, |
|
|
**kwargs, |
|
|
): |
|
|
hidden_states = self.embed_tokens(input_ids) |
|
|
|
|
|
for layer in self.layers: |
|
|
hidden_states = layer(hidden_states) |
|
|
|
|
|
hidden_states = self.norm(hidden_states) |
|
|
|
|
|
return BaseModelOutputWithPast( |
|
|
last_hidden_state=hidden_states, |
|
|
hidden_states=None, |
|
|
past_key_values=None, |
|
|
) |
|
|
|
|
|
class ThetaForCausalLM(PreTrainedModel): |
|
|
config_class = ThetaConfig |
|
|
|
|
|
def __init__(self, config: ThetaConfig): |
|
|
super().__init__(config) |
|
|
self.model = ThetaModel(config) |
|
|
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
input_ids=None, |
|
|
labels=None, |
|
|
**kwargs, |
|
|
): |
|
|
outputs = self.model(input_ids=input_ids, **kwargs) |
|
|
logits = self.lm_head(outputs.last_hidden_state) |
|
|
|
|
|
loss = None |
|
|
if labels is not None: |
|
|
shift_logits = logits[..., :-1, :].contiguous() |
|
|
shift_labels = labels[..., 1:].contiguous() |
|
|
loss_fct = nn.CrossEntropyLoss() |
|
|
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) |
|
|
|
|
|
return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits} |
|
|
|
