| 185 | |
| 186 | |
| 187 | class LanguageModel(nn.Module): |
| 188 | def __init__(self, config: TextConfig): |
| 189 | super().__init__() |
| 190 | self.model_type = config.model_type |
| 191 | if self.model_type != "llama": |
| 192 | raise ValueError( |
| 193 | f"Model type {self.model_type} not supported. Currently only 'llama' is supported" |
| 194 | ) |
| 195 | self.model = Llama(config) |
| 196 | self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| 197 | |
| 198 | def __call__( |
| 199 | self, |
| 200 | inputs: mx.array, |
| 201 | cache=None, |
| 202 | inputs_embeds=None, |
| 203 | ): |
| 204 | out, cache = self.model(inputs, cache, inputs_embeds) |
| 205 | return self.lm_head(out), cache |
| 206 | |
| 207 | @staticmethod |
| 208 | def sanitize(weights): |
| 209 | # Remove unused precomputed rotary freqs |
| 210 | return { |
| 211 | k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k |
| 212 | } |