huggingfaceのgenerationの関数をtorch modelから使えるようにしたい

def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
    top_k = min(self.top_k, scores.size(-1))  # Safety check
    # Remove all tokens with a probability less than the last token of the top-k
    indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
    scores = scores.masked_fill(indices_to_remove, self.filter_value)
    return scores

next_token_scores = logits_processor(input_ids, next_token_logits)
next_token_scores = logits_warper(input_ids, next_token_scores)

...

probs = nn.functional.softmax(next_token_scores, dim=-1)
next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)

next_tokens = torch.argmax(next_tokens_scores, dim=-1)

from transformers.generation.utils import (
    LogitsProcessorList, 
    TopKLogitsWarper,
    TopPLogitsWarper,
    TemperatureLogitsWarper,
    RepetitionPenaltyLogitsProcessor
)

logits_processor = LogitsProcessorList([
    RepetitionPenaltyLogitsProcessor(repetition_penalty),
])

logits_wraper = LogitsProcessorList([
        TemperatureLogitsWarper(temperature),
        TopPLogitsWarper(top_p),
        TopKLogitsWarper(top_k),
])

## modelからlogitsを得る
logits = model(x, input_pos)

logits = logits[:, -1, :] ## sequenceの最後のlogit
next_token_scores = logits_processor(x, logits)
next_token_scores = logits_wraper(x, next_token_scores)
next_token_scores = next_token_scores.squeeze(0) ## 次元の調整

## next_token_scoresに対してsoftmaxで確率にして、次のtokenのindexを出力する
probs = torch.nn.functional.softmax(next_token_scores, dim=-1)
next_idx = torch.multinomial(probs, num_samples=1)