|
|
| import torch |
| import torch.nn.functional as F |
|
|
| |
| class CustomRepetitionPenaltyLogitsProcessorRepeat(): |
|
|
| def __init__(self, penalty: float, max_input_ids, past_window): |
| if not isinstance(penalty, float) or not (penalty > 0): |
| raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") |
|
|
| self.penalty = penalty |
| self.max_input_ids = max_input_ids |
| self.past_window = past_window |
|
|
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: |
| |
| input_ids = input_ids[:, -self.past_window:] |
| freq = F.one_hot(input_ids, scores.size(1)).sum(1) |
| freq[self.max_input_ids:] = 0 |
| alpha = self.penalty**freq |
| scores = torch.where(scores < 0, scores*alpha, scores/alpha) |
|
|
| return scores |
| |
| class CustomRepetitionPenaltyLogitsProcessor(): |
|
|
| def __init__(self, penalty: float, max_input_ids, past_window): |
| if not isinstance(penalty, float) or not (penalty > 0): |
| raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") |
|
|
| self.penalty = penalty |
| self.max_input_ids = max_input_ids |
| self.past_window = past_window |
|
|
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: |
| |
| input_ids = input_ids[:, -self.past_window:] |
| score = torch.gather(scores, 1, input_ids) |
| _score = score.detach().clone() |
| score = torch.where(score < 0, score * self.penalty, score / self.penalty) |
| score[input_ids>=self.max_input_ids] = _score[input_ids>=self.max_input_ids] |
| scores.scatter_(1, input_ids, score) |
| |
| return scores |
