大语言模型之生成/推理：参数与解码策略原理及其代码实现

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "llama-2-7b-hf" # 用你的模型的地址
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "say"
inputs = tokenizer(text, return_tensors="pt")
print(f"inputs:{inputs}")

# 结果
inputs:{'input_ids': tensor([[   1, 1827]]), 'attention_mask': tensor([[1, 1]])}

logits = model.forward(input_ids)
print("Logits Shape:", logits.logits.shape)
print(f"logits:{logits.logits}")

# 结果
Logits Shape: torch.Size([1, 2, 32000])
logits:tensor([[[-12.9696,  -7.4491,  -0.4354,  ...,  -6.8250,  -8.0804,  -7.5782],
         [-11.3775, -10.1338,  -2.3563,  ...,  -6.7709,  -6.5252,  -8.9753]]],
       device='cuda:0', grad_fn=<UnsafeViewBackward0>)

# 在最后一个维度上（32000）进行求最大值操作，返回具有最高分数的词汇索引
next_token = torch.argmax(logits.logits, dim=-1).reshape(-1)[1]
print(f"next_token:{next_token}")

# 结果
next_token:tensor([22172], device='cuda:0')

next_word = tokenizer.decode(next_token)
print(f"next_word:{next_word}")

# 结果
next_word:hello

import torch
logits = torch.tensor([[0.5, 1.2, -1.0, 0.1]])
# 无temperature
probs = torch.softmax(logits, dim=-1)
# temperature low 0.5
probs_low = torch.softmax(logits / 0.5, dim=-1)
# temperature high 2
probs_high = torch.softmax(logits / 2, dim=-1)

print(f"probs:{probs}")
print(f"probs_low:{probs_low}")
print(f"probs_high:{probs_high}")
# 结果
probs: tensor([[0.2559, 0.5154, 0.0571, 0.1716]])
probs_low: tensor([[0.1800, 0.7301, 0.0090, 0.0809]])
probs_high: tensor([[0.2695, 0.3825, 0.1273, 0.2207]])

import torch
# 样例：probs: tensor([[0.2559, 0.5154, 0.0571, 0.1716]])
probs = torch.tensor([[0.2559, 0.5154, 0.0571, 0.1716]])
# 第一步进行排序
probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
# 结果
probs_sort: tensor([[0.5154, 0.2559, 0.1716, 0.0571]])
probs_idx: tensor([[1, 0, 3, 2]])

# 第二步概率的累积和
probs_sum = torch.cumsum(probs_sort, dim=-1)
# 结果
probs_sum: tensor([[0.5154, 0.7713, 0.9429, 1.0000]])

# 第三步找到第一个大于阈值p的位置，假设p=0.9，并将后面的概率值置为0：
mask = probs_sum - probs_sort > p
probs_sort[mask] = 0.0
# 结果
probs_sort: tensor([[0.5154, 0.2559, 0.1716, 0.0000]])

# 第四步复原原序列
new_probs = probs_sort.scatter(1, probs_idx, probs_sort)
# 结果
new_probs: tensor([[0.2559, 0.5154, 0.0000, 0.1716]])

# 注：在真实实现中一般会把舍弃的概率置为-inf，即
zero_indices = (new_probs == 0)
new_probs[zero_indices] = float('-inf')
# 结果
new_probs: tensor([[0.2559, 0.5154, -inf, 0.1716]])
# 完整代码
def sample_top_p(probs, p):
    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    mask = probs_sum - probs_sort > p
    probs_sort[mask] = 0.0
    new_probs = probs_sort.scatter(1, probs_idx, probs_sort)
    zero_indices = (new_probs == 0)
    new_probs[zero_indices] = float('-inf')
    return new_probs

import torch
filter_value = -float("Inf")
top_k = 2
probs = torch.tensor([[0.2559, 0.5154, 0.0571, 0.1716]])
indices_to_remove = probs < torch.topk(probs, top_k)[0][..., -1, None]
new_probs = probs.masked_fill(indices_to_remove, filter_value)
print("new_probs:", new_probs)
# 结果
new_probs: tensor([[0.2559, 0.5154,   -inf,   -inf]])

import numpy as np
def apply_repetition_penalty(probs, repetition_penalty, prev_tokens):
    adjusted_probs = np.copy(probs)
    for token in set(prev_tokens):
        adjusted_probs[token] *= repetition_penalty
    adjusted_probs /= np.sum(adjusted_probs)  
    return adjusted_probs
# 示例概率分布，索引对应不同词语
original_probs = np.array([0.3, 0.1, 0.3, 0.1, 0.2])
# 示例先前生成的词语
previous_tokens = [2, 4, 2]
# 重复惩罚系数
repetition_penalty = 0.8
# 应用重复惩罚，得到调整后的概率分布
adjusted_probs = apply_repetition_penalty(original_probs, repetition_penalty, previous_tokens)

print("原始概率分布：", original_probs)
print("调整后的概率分布：", adjusted_probs)

# 结果
原始概率分布： [0.3 0.1 0.3 0.1 0.2]
调整后的概率分布： [0.33333333 0.11111111 0.26666667 0.11111111 0.17777778]

import torch
probs = torch.tensor([[0.2559, 0.5154, 0.0571, 0.1716]])
next_token = torch.multinomial(probs, num_samples=1)
print("next_token:", next_token)
# 结果
next_token: tensor([[1]])

UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.5` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.
UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.1` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`.
UserWarning: `do_sample` is set to `False`. However, `top_k` is set to `20` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_k`.

class BeamSearchNode:
    def __init__(self, sequence, score):
        self.sequence = sequence  # 生成的序列
        self.score = score  # 分数（概率）

# 示例：下一个token的概率函数，简单使用固定概率
def simple_next_word_probs(sequence):
    if sequence[-1] == "<end>":
        return {}
    return {"apple": 0.3, "like": 0.35, "peach": 0.2, "banana": 0.15}

def beam_search(initial_sequence, next_word_probs_func, num_beams, max_sequence_length):
    # 初始化初始节点，且分数为1
    initial_node = BeamSearchNode(sequence=initial_sequence, score=1.0)
    candidates = [initial_node]

    final_candidates = []  # 最终的候选序列
    # 只要候选节点列表不为空，且 final_candidates 中的候选节点数量还没有达到指定的束宽度，就继续进行搜索
    while candidates and len(final_candidates) < num_beams:
        # 候选节点排序
        candidates.sort(key=lambda x: -x.score)
        current_node = candidates.pop(0)
        # 当节点序列末尾生成结束符号（如"<end>"），或者当生成的序列长度达到最大限制时终止节点的扩展
        if current_node.sequence[-1] == "<end>" or len(current_node.sequence) >= max_sequence_length:
            final_candidates.append(current_node)
        else:
            # 获取下一个token的概率，我们的例子返回的是固定的概率
            next_words_probs = next_word_probs_func(current_node.sequence) 
            # 生成新的候选序列，并计算分数           
            for next_word, next_word_prob in next_words_probs.items():
                new_sequence = current_node.sequence + [next_word]
                new_score = current_node.score * next_word_prob
                new_node = BeamSearchNode(sequence=new_sequence, score=new_score)
                candidates.append(new_node)

    return [candidate.sequence for candidate in final_candidates]

initial_sequence = ["<start>", "I"]
num_beams = 3
max_sequence_length = 3
result = beam_search(initial_sequence, simple_next_word_probs, num_beams, max_sequence_length)

for idx, sequence in enumerate(result):
    print(f"Sentence {idx + 1}: {' '.join(sequence)}")

# 结果，符合我们的概率分布
Sentence 1: <start> I like
Sentence 2: <start> I apple
Sentence 3: <start> I peach

# 当长度为4时呢？
Sentence 1: <start> I like like
Sentence 2: <start> I like apple
Sentence 3: <start> I apple like

# 为什么结果变了呢？我们来看一下最终每个序列的分数
current_node: ['<start>', 'I', 'like', 'like']
current_score: 0.12249999999999998
current_node: ['<start>', 'I', 'like', 'apple']
current_score: 0.105
current_node: ['<start>', 'I', 'apple', 'like']
current_score: 0.105

# 再看一下其他序列的分数，apple的
new_node: ['<start>', 'I', 'apple', 'apple']
new_score: 0.09
new_node: ['<start>', 'I', 'apple', 'like']
new_score: 0.105
new_node: ['<start>', 'I', 'apple', 'peach']
new_score: 0.06
new_node: ['<start>', 'I', 'apple', 'banana']
new_score: 0.045

# 再看一下其他序列的分数，peach的
new_node: ['<start>', 'I', 'peach', 'apple']
new_score: 0.06
new_node: ['<start>', 'I', 'peach', 'like']
new_score: 0.06999999999999999
new_node: ['<start>', 'I', 'peach', 'peach']
new_score: 0.04000000000000001
new_node: ['<start>', 'I', 'peach', 'banana']
new_score: 0.03

"`num_beams` needs to be greater than 1 for constrained generation."
"`do_sample` needs to be false for constrained generation."
"`num_beam_groups` not supported yet for constrained generation."

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
import torch
model_name = "llama-2-7b-hf" # 你模型的位置
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "say hello to"
inputs = tokenizer(text, return_tensors="pt")
print(f"inputs:{inputs}")
input_ids = inputs["input_ids"].to("cuda")

# generate实现
generation_output = model.generate(
    input_ids=input_ids,
    num_beams = 3,
    num_return_sequences=3,
    return_dict_in_generate=True,
    max_new_tokens=3,
)

print("query:", text)
for i, output_sequence in enumerate(generation_output.sequences):
    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)
    print(f"Generated sequence {i+1}: {output_text}")

# 结果
inputs:{'input_ids': tensor([[    1,  1827, 22172,   304]]), 'attention_mask': tensor([[1, 1, 1, 1]])}
query: say hello to
Generated sequence 1: say hello to your new favorite
Generated sequence 2: say hello to your new best
Generated sequence 3: say hello to our newest

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
import torch
model_name = "llama-2-7b-hf" # 你模型的位置
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "say hello to"
inputs = tokenizer(text, return_tensors="pt")
print(f"inputs:{inputs}")
input_ids = inputs["input_ids"].to("cuda")

force_words = ["my"]
force_words_ids = tokenizer(force_words, add_special_tokens=False).input_ids

generation_output = model.generate(
    input_ids=input_ids,
    force_words_ids = force_words_ids,
    num_beams = 3,
    num_return_sequences=3,
    return_dict_in_generate=True,
    max_new_tokens=3,
)

print("query:", text)
for i, output_sequence in enumerate(generation_output.sequences):
    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)
    print(f"Generated sequence {i+1}: {output_text}")

# 结果    
inputs:{'input_ids': tensor([[    1,  1827, 22172,   304]]), 'attention_mask': tensor([[1, 1, 1, 1]])}
query: say hello to
Generated sequence 1: say hello to my little friend
Generated sequence 2: say hello to your new favorite
Generated sequence 3: say hello to your new my

def ranking(context_hidden, next_hidden, next_top_k_ids, next_top_k_probs, alpha):
    '''
       该函数是实现Contrastive Search中next token预测中候选token的排序分数，分数最大对应token为输出结果
        context_hidden: beam_width x context_len x embed_dim ,用于计算相似度，是公式中x_j集合表征向量
        next_hidden: beam_width x 1 x embed_dim，用于计算相似度，是公式中候选token v 的表征向量
        next_top_k_ids: beam_width x 1，记录候选token的编码
        next_top_k_probs，候选token的模型预测概率
        alpha，惩罚参数
    '''
    beam_width, context_len, embed_dim = context_hidden.size()
    assert next_hidden.size() == torch.Size([beam_width, 1, embed_dim])
    norm_context_hidden = context_hidden / context_hidden.norm(dim=2, keepdim=True) 
    norm_next_hidden = next_hidden / next_hidden.norm(dim=2, keepdim=True)
    cosine_matrix = torch.matmul(norm_context_hidden, norm_next_hidden.transpose(1,2)).squeeze(-1) #计算相似度矩阵
    assert cosine_matrix.size() == torch.Size([beam_width, context_len])
    scores, _ = torch.max(cosine_matrix, dim = -1) #输出公式第二项值
    assert scores.size() == torch.Size([beam_width])
    next_top_k_probs = next_top_k_probs.view(-1)  #输出公式第一项值
    scores = (1.0 - alpha) * next_top_k_probs - alpha * scores  #对应公式整体计算
    _, selected_idx = torch.topk(scores, k = 1)
    assert selected_idx.size() == torch.Size([1])
    selected_idx = selected_idx.unsqueeze(0)
    assert selected_idx.size() == torch.Size([1,1])
    next_id = torch.gather(next_top_k_ids, dim = 0, index=selected_idx)
    assert next_id.size() == torch.Size([1,1])
    return next_id

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
import torch
model_name = "llama-2-7b-hf" # 你模型的位置
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)


text = "say hello to"
inputs = tokenizer(text, return_tensors="pt")
print(f"inputs:{inputs}")
input_ids = inputs["input_ids"].to("cuda")


generation_output = model.generate(
    input_ids=input_ids,
    penalty_alpha = 0.5,
    top_k = 30,
    return_dict_in_generate=True,
    max_new_tokens=3,
)

# 直接使用其函数
# generation_output = model.contrastive_search(
#     input_ids=input_ids,
#     penalty_alpha = 0.5,
#     top_k = 30,
#     return_dict_in_generate=True,
#     max_new_tokens=3,
# )

print("query:", text)
for i, output_sequence in enumerate(generation_output.sequences):
    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)
    print(f"Generated sequence {i+1}: {output_text}")

# 结果
inputs:{'input_ids': tensor([[    1,  1827, 22172,   304]]), 'attention_mask': tensor([[1, 1, 1, 1]])}
query: say hello to
Generated sequence 1: say hello to 20

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
import torch
model_name = "llama-2-7b-hf" # 你模型的位置
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "say hello to"
inputs = tokenizer(text, return_tensors="pt")
print(f"inputs:{inputs}")
input_ids = inputs["input_ids"].to("cuda")


generation_output = model.generate(
    input_ids=input_ids,
    num_beams = 1,
    do_sample = False,
    return_dict_in_generate=True,
    max_new_tokens=3,
)
# 直接指定使用其函数
# generation_output = model.greedy_search(
#     input_ids=input_ids,
#     num_beams = 1,
#     do_sample = False,
#     return_dict_in_generate=True,
#     max_length = 7
# )


print("query:", text)
for i, output_sequence in enumerate(generation_output.sequences):
    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)
    print(f"Generated sequence {i+1}: {output_text}")

# 结果
inputs:{'input_ids': tensor([[    1,  1827, 22172,   304]]), 'attention_mask': tensor([[1, 1, 1, 1]])}
query: say hello to
Generated sequence 1: say hello to the newest

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from transformers import (
    LogitsProcessorList,
    TopKLogitsWarper,
    TopPLogitsWarper,
    TemperatureLogitsWarper,
    )

import torch
model_name = "llama-2-7b-hf" # 你模型的位置
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "say hello to"
inputs = tokenizer(text, return_tensors="pt")
print(f"inputs:{inputs}")
input_ids = inputs["input_ids"].to("cuda")


generation_output = model.generate(
    input_ids=input_ids,
    num_beams = 1,
    do_sample = True,
    temperature = 1.2,
    top_k = 100,
    top_p = 0.6,
    return_dict_in_generate=True,
    max_length=7,
)


# sample实现
# logits_warper = LogitsProcessorList(
#     [
#     TopKLogitsWarper(100),
#     TemperatureLogitsWarper(1.2),
#     TopPLogitsWarper(0.6)
#     ]
# )
# generation_output = model.sample(
#     input_ids=input_ids,
#     logits_warper=logits_warper,
#     return_dict_in_generate=True,
#     max_length=7,
# )


print("query:", text)
for i, output_sequence in enumerate(generation_output.sequences):
    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)
    print(f"Generated sequence {i+1}: {output_text}")

# 注意这种方式每次结果都可能不一样
inputs:{'input_ids': tensor([[    1,  1827, 22172,   304]]), 'attention_mask': tensor([[1, 1, 1, 1]])}
query: say hello to
Generated sequence 1: say hello to our new intern

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
import torch
model_name = "llama-2-13b-hf" # 你自己模型的位置
assistant_model_name = "llama-2-7b-hf" # 你自己模型的位置
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_model_name, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "say hello to"
inputs = tokenizer(text, return_tensors="pt")
print(f"inputs:{inputs}")
input_ids = inputs["input_ids"].to("cuda")


generation_output = model.generate(
    assistant_model=assistant_model,
    input_ids=input_ids,
    num_beams = 1,
    do_sample = False,
    return_dict_in_generate=True,
    max_length=7,
)


print("query:", text)
for i, output_sequence in enumerate(generation_output.sequences):
    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)
    print(f"Generated sequence {i+1}: {output_text}")

# 结果
inputs:{'input_ids': tensor([[    1,  1827, 22172,   304]]), 'attention_mask': tensor([[1, 1, 1, 1]])}
query: say hello to
Generated sequence 1: say hello to the newest