MCPcopy Create free account
hub / github.com/OpenSparseLLMs/MoM / forward

Method forward

mom/models/mom/modeling_mom.py:264–350  ·  view source on GitHub ↗
(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,  # noqa
        inputs_embeds: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs: Unpack[Dict]
    )

Source from the content-addressed store, hash-verified

262 self.embeddings = value
263
264 def forward(
265 self,
266 input_ids: Optional[torch.LongTensor] = None,
267 attention_mask: Optional[torch.Tensor] = None, # noqa
268 inputs_embeds: Optional[torch.FloatTensor] = None,
269 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
270 use_cache: Optional[bool] = None,
271 output_attentions: Optional[bool] = None,
272 output_hidden_states: Optional[bool] = None,
273 return_dict: Optional[bool] = None,
274 **kwargs: Unpack[Dict]
275 ) -> Union[Tuple, BaseModelOutputWithPast]:
276 if output_attentions:
277 warnings.warn("`MomModel` does not `output_attentions` now, setting it to `False`.")
278 output_attentions = False
279 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
280 output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
281 use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
282 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
283
284 # retrieve input_ids and inputs_embeds
285 if input_ids is not None and inputs_embeds is not None:
286 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
287 if input_ids is None and inputs_embeds is None:
288 raise ValueError("You have to specify either input_ids or inputs_embeds")
289
290 # import debugpy
291 # debugpy.connect(5678)
292
293 if inputs_embeds is None:
294 inputs_embeds = self.embeddings(input_ids)
295 hidden_states = inputs_embeds
296
297 if use_cache and not isinstance(past_key_values, Cache):
298 past_key_values = Cache.from_legacy_cache(past_key_values)
299
300 if self.gradient_checkpointing and self.training and use_cache:
301 logger.warning_once("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...")
302 use_cache = False
303
304 all_hidden_states = () if output_hidden_states else None
305 all_attns = () if output_attentions else None
306 all_router_logits = ()
307
308 for layer in self.layers:
309 if output_hidden_states:
310 all_hidden_states += (hidden_states,)
311
312 if self.gradient_checkpointing and self.training:
313 hidden_states, attentions, past_key_values, router_logits = self._gradient_checkpointing_func(
314 layer.__call__,
315 hidden_states,
316 attention_mask,
317 past_key_values,
318 use_cache,
319 output_attentions,
320 **kwargs
321 )

Callers

nothing calls this directly

Calls 1

MomOutputWithPastClass · 0.85

Tested by

no test coverage detected