1
2
3
4
5
6

if self.prompt_ignore_length:
    input_ids = input_ids[:, self.prompt_ignore_length :]
score = torch.gather(scores, 1, input_ids)

score = torch.where(score < 0, score * self.penalty, score / self.penalty)
scores_processed = scores.scatter(1, input_ids, score)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.logits_process import RepetitionPenaltyLogitsProcessor

def print_top_tokens(tokenizer, scores, message=""):
    
    probs = F.softmax(scores, dim=-1)
    top_scores = torch.topk(scores[0], 5)

    print(f"\n{message}")
    print("-" * 50)
    print(f"{'Token':<15} {'Raw Score':<15} {'Probability':<15}")
    print("-" * 50)

    for idx, (score, prob) in enumerate(
        zip(top_scores.values, probs[0][top_scores.indices])
    ):
        token = tokenizer.decode([top_scores.indices[idx]])
        print(f"{token:<15} {score.item():>8.3f}       {prob.item():>8.6f}")


model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")


text = "I love coding. I love"
input_ids = tokenizer.encode(text, return_tensors="pt")


with torch.no_grad():
    outputs = model(input_ids)
    original_scores = outputs.logits[:, -1, :].clone()  

    
    penalty_values = [0.8, 1.2, 2.0]

    print(f"输入文本: {text}")

    
    print_top_tokens(tokenizer, original_scores, "原始输出 (无repetition penalty)")

    
    for penalty in penalty_values:
        processor = RepetitionPenaltyLogitsProcessor(penalty=penalty)
        processed_scores = processor(input_ids, original_scores.clone())
        print_top_tokens(
            tokenizer, processed_scores, f"应用repetition penalty = {penalty}后的输出"
        )

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

输入文本: I love coding. I love
原始输出 (无repetition penalty)
--------------------------------------------------
Token           Raw Score       Probability
--------------------------------------------------
 the              16.583       0.176431
 to               15.963       0.094929
 learning         15.550       0.062831
 solving          15.482       0.058693
 programming      15.221       0.045199

应用repetition penalty = 0.8后的输出
--------------------------------------------------
Token           Raw Score       Probability
--------------------------------------------------
 coding           18.377       0.519966
 the              16.583       0.086431
 to               15.963       0.046504
 learning         15.550       0.030780
 solving          15.482       0.028753

应用repetition penalty = 1.2后的输出
--------------------------------------------------
Token           Raw Score       Probability
--------------------------------------------------
 the              16.583       0.180972
 to               15.963       0.097372
 learning         15.550       0.064449
 solving          15.482       0.060203
 programming      15.221       0.046362

应用repetition penalty = 2.0后的输出
--------------------------------------------------
Token           Raw Score       Probability
--------------------------------------------------
 the              16.583       0.181423
 to               15.963       0.097615
 learning         15.550       0.064609
 solving          15.482       0.060353
 programming      15.221       0.046477