(self, config: TextConfig)
| 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, |