File size: 1,107 Bytes
2811206 | 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 | import torch
from torch import nn
from transformers.modeling_utils import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput
class Phi3ForCausalLM(PreTrainedModel):
config_class = Phi3Config
base_model_prefix = "phi3"
def __init__(self, config):
super().__init__(config)
self.hidden_size = config.hidden_size
self.num_hidden_layers = config.num_hidden_layers
self.num_attention_heads = config.num_attention_heads
self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
self.layers = nn.ModuleList([nn.TransformerEncoderLayer(config.hidden_size, config.num_attention_heads) for _ in range(config.num_hidden_layers)])
self.output_layer = nn.Linear(config.hidden_size, config.vocab_size)
def forward(self, input_ids):
embeddings = self.embedding(input_ids)
hidden_states = embeddings
for layer in self.layers:
hidden_states = layer(hidden_states)
logits = self.output_layer(hidden_states)
return BaseModelOutput(last_hidden_state=logits)
|