59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from transformers.generation import TopKLogitsWarper, TopPLogitsWarper
|
|
|
|
|
|
class CustomRepetitionPenaltyLogitsProcessorRepeat:
|
|
|
|
def __init__(self, penalty: float, max_input_ids: int, past_window: int):
|
|
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:
|
|
if input_ids.size(1) > self.past_window:
|
|
input_ids = input_ids.narrow(1, -self.past_window, self.past_window)
|
|
freq = F.one_hot(input_ids, scores.size(1)).sum(1)
|
|
if freq.size(0) > self.max_input_ids:
|
|
freq.narrow(
|
|
0, self.max_input_ids, freq.size(0) - self.max_input_ids
|
|
).zero_()
|
|
alpha = torch.pow(self.penalty, freq)
|
|
scores = scores.contiguous()
|
|
inp = scores.multiply(alpha)
|
|
oth = scores.divide(alpha)
|
|
con = scores < 0
|
|
out = torch.where(con, inp, oth)
|
|
del inp, oth, scores, con, alpha
|
|
return out
|
|
|
|
|
|
def gen_logits(
|
|
num_code: int,
|
|
top_P=0.7,
|
|
top_K=20,
|
|
repetition_penalty=1.0,
|
|
):
|
|
logits_warpers = []
|
|
if top_P is not None:
|
|
logits_warpers.append(TopPLogitsWarper(top_P, min_tokens_to_keep=3))
|
|
if top_K is not None:
|
|
logits_warpers.append(TopKLogitsWarper(top_K, min_tokens_to_keep=3))
|
|
|
|
logits_processors = []
|
|
if repetition_penalty is not None and repetition_penalty != 1:
|
|
logits_processors.append(
|
|
CustomRepetitionPenaltyLogitsProcessorRepeat(
|
|
repetition_penalty, num_code, 16
|
|
)
|
|
)
|
|
|
|
return logits_warpers, logits_processors
|