different sampling strategies. top_p, typical, ...

[fabianwinter93]

It would be easy to implement all sorts of sampling methods by first calling _get_logits of the parent model class via super() and then adding the processing. Should i try to contribute this or is this not the way to go?

class SamplingModel(g.models.MistralChat):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.top_p = 1. if 'top_p' not in kwargs else kwargs['top_p']

  def _get_logits(self, token_ids, forced_bytes, current_temp):
    logits = super()._get_logits(token_ids, forced_bytes, current_temp)

    if self.top_p < 1.:
      logits = torch.tensor(logits)

      sorted_logits, sorted_indices = torch.sort(logits, descending=True)
      cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)

      sorted_indices_to_remove = cumulative_probs > self.top_p

      sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
      sorted_indices_to_remove[..., 0] = 0

      indices_to_remove = sorted_indices[sorted_indices_to_remove]
      logits[indices_to_remove] = -float('inf')

      self._cache_state["logits"] = logits

    return logits