(self, config: CohereConfig)
| 92 | class CohereModel(Module): |
| 93 | |
| 94 | def __init__(self, config: CohereConfig) -> None: |
| 95 | super().__init__() |
| 96 | |
| 97 | self.mapping = config.mapping |
| 98 | if self.mapping.is_first_pp_rank(): |
| 99 | self.vocab_embedding = Embedding(config.vocab_size, |
| 100 | config.hidden_size, |
| 101 | dtype=config.dtype, |
| 102 | tp_group=config.mapping.tp_group, |
| 103 | tp_size=config.mapping.tp_size, |
| 104 | tp_rank=config.mapping.tp_rank) |
| 105 | |
| 106 | self.layers = DecoderLayerList(CohereDecoderLayer, config) |
| 107 | |
| 108 | if self.mapping.is_last_pp_rank(): |
| 109 | self.ln_f = LayerNorm(normalized_shape=config.hidden_size, |
| 110 | eps=config.norm_epsilon, |
| 111 | bias=False, |
| 112 | dtype=config.dtype) |
| 113 | |
| 114 | def forward( |
| 115 | self, |
no test coverage detected