Method infer_

moss_inference.py:315–340

Inference method that computes logits and past key values. Args: input_ids (torch.Tensor): The input IDs tensor. attention_mask (torch.Tensor): The attention mask tensor. past_key_values (Optional[Tuple[torch.Tensor]]): The past key values tuple.

(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        past_key_values: Optional[Tuple[torch.Tensor]],
    )

        return logits
    
    def infer_(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        past_key_values: Optional[Tuple[torch.Tensor]],
    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
        """
        Inference method that computes logits and past key values.

        Args:
            input_ids (torch.Tensor): The input IDs tensor.
            attention_mask (torch.Tensor): The attention mask tensor.
            past_key_values (Optional[Tuple[torch.Tensor]]): The past key values tuple.

        Returns:
            Tuple[torch.Tensor, Tuple[torch.Tensor]]: A tuple containing the logits and past key values.
        """
        inputs = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
        }
        with torch.no_grad():
            outputs: BaseModelOutputWithPast = self.model(**inputs)

        return outputs.logits, outputs.past_key_values

    def __call__(self, input):
        return self.forward(input)