(
self,
input_ids=None,
position_ids=None,
use_cache=False,
attention_mask=None,
spec_decoding_params=None,
kv_cache_params=None,
attention_params=None,
hidden_states=None,
)
| 112 | dtype=config.dtype) |
| 113 | |
| 114 | def forward( |
| 115 | self, |
| 116 | input_ids=None, |
| 117 | position_ids=None, |
| 118 | use_cache=False, |
| 119 | attention_mask=None, |
| 120 | spec_decoding_params=None, |
| 121 | kv_cache_params=None, |
| 122 | attention_params=None, |
| 123 | hidden_states=None, |
| 124 | ): |
| 125 | if self.mapping.is_first_pp_rank(): |
| 126 | hidden_states = self.vocab_embedding(input_ids) |
| 127 | else: |
| 128 | hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) |
| 129 | |
| 130 | hidden_states = self.layers.forward( |
| 131 | hidden_states, |
| 132 | use_cache=use_cache, |
| 133 | attention_mask=attention_mask, |
| 134 | kv_cache_params=kv_cache_params, |
| 135 | attention_params=attention_params, |
| 136 | spec_decoding_params=spec_decoding_params) |
| 137 | |
| 138 | if use_cache: |
| 139 | hidden_states, presents = hidden_states |
| 140 | |
| 141 | if self.mapping.is_last_pp_rank(): |
| 142 | hidden_states = self.ln_f(hidden_states) |
| 143 | else: |
| 144 | hidden_states = send(hidden_states, self.mapping.next_pp_rank()) |
| 145 | |
| 146 | if use_cache: |
| 147 | return (hidden_states, presents) |
| 148 | return hidden_states |
| 149 | |
| 150 | |
| 151 | class CohereForCausalLM(DecoderModelForCausalLM): |
nothing calls this directly
no test coverage detected