| 237 | |
| 238 | class MossModel(Module): |
| 239 | def __init__(self, config): |
| 240 | super(MossModel, self).__init__() |
| 241 | |
| 242 | self.config = config |
| 243 | self.embed_dim = config.n_embd |
| 244 | self.vocab_size = config.vocab_size |
| 245 | self.wte = nn.Embedding(config.vocab_size, self.embed_dim) |
| 246 | self.drop = nn.Dropout(config.embd_pdrop) |
| 247 | self.h = nn.ModuleList([MossBlock(config) for _ in range(config.n_layer)]) |
| 248 | self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) |
| 249 | self.rotary_dim = min(config.rotary_dim, config.n_ctx // config.n_head) |
| 250 | |
| 251 | self.gradient_checkpointing = False |
| 252 | |
| 253 | self.apply(partial(_init_weights, config)) |
| 254 | |
| 255 | def execute( |
| 256 | self, |