这段代码定义了一个简单的对话生成系统,包括模型加载、词汇表加载、以及基于给定提示生成文本的功能。下面是对代码的解析:
-
load_model_and_voc(device="cpu")
:- 该函数用于加载预训练的模型和词汇表(vocabulary)。它首先从文件
total_voc.pkl
中加载词汇表,并创建一个名为SamOut
的神经网络实例。 - 模型参数的数量被打印出来以供参考。
- 然后尝试加载指定路径下的预训练权重到模型中,并将模型移动到指定的设备(CPU 或 GPU)上。
- 最后设置模型为评估模式(
.eval()
),并返回模型和词汇表。
- 该函数用于加载预训练的模型和词汇表(vocabulary)。它首先从文件
-
gen_token(voc, model, prompt, max_len, rp=1.2, temp=0.13, top_k=16, device="cpu")
:- 这个函数负责根据提供的提示(prompt)生成新的文本序列。
- 它接受多个参数,包括词汇表、模型、初始提示、最大生成长度等。
- 函数内部实现了重复抑制、温度调整和top-k采样等技术来控制生成文本的质量。
- 使用softmax函数对模型输出进行处理,并通过多类别抽样选择下一个token。
- 如果生成了特殊的开始标记
<|sos|>
,则停止生成过程。 - 生成的每个token会立即打印在屏幕上,形成即时响应的效果。
-
t_infre()
:- 此函数是交互式推理循环,允许用户输入文本,然后调用
gen_token
函数来生成回应。 - 它是一个无限循环,持续等待用户的输入直到程序被手动终止。
- 此函数是交互式推理循环,允许用户输入文本,然后调用
-
if __name__ == '__main__':
- 这部分代码确保当脚本作为主程序运行时,会执行某些特定的操作或测试。
- 注释掉的代码可能是之前用于数据预处理、训练或其他实验的部分。
- 最终调用了
t_infre()
函数来启动交互式推理。
需要注意的是,这里使用的 SamOut
类并没有在给出的代码片段中定义,因此你可能需要确保这个类已经被正确实现并在其他地方导入。此外,为了使代码能够正常工作,你需要确保所有依赖库(如 PyTorch 和 pandas)已经安装,并且所有提及的数据文件和模型权重文件都存在于正确的路径下。
def load_model_and_voc(device="cpu"):voc = pd.read_pickle("total_voc.pkl")net = SamOut(len(voc["voc"]), 1024 + 512, 64, 16)# net = SamOut(len(voc["voc"]), 512, 32, 8)print(sum([i.shape[0] * i.shape[1] for i in net.parameters() if len(i.shape) > 1]) + sum([i.shape[0] for i in net.parameters() if len(i.shape) == 1]))# net.load_state_dict(torch.load("pretrain_768.pth", map_location=device))# net.load_state_dict(torch.load("pretrain_sft_single.pth", map_location=device))net.load_state_dict(torch.load("pretrain_sft_single_1024.pth", map_location=device))# net.load_state_dict(torch.load("pretrain.pth", map_location=device))net.to(device)net.eval()return net, vocdef gen_token(voc, model, prompt, max_len, rp=1.2, temp=0.13, top_k=16, device="cpu"):print("agent:", end="", flush=True)for _ in range(max_len):prompt_list = []for i in prompt:if i not in voc["voc"]:prompt_list += [voc["voc"].index(ii) for ii in voc["voc0"].get(i)]else:prompt_list.append(voc["voc"].index(i))out, _ = model(torch.Tensor([prompt_list]).to(device).long())out = out[:, -1:]# 重复抑制for token_id in enumerate(prompt_list):out[:, :, token_id] /= rpscore = torch.softmax(out, -1)[0, 0]score, score_index = torch.sort(score,descending=True)score=score.detach().numpy()score_sum = np.cumsum(score)score_index = score_index.detach().numpy()score1=score[score_sum<0.8]if score1.size==0:score=score[:1]else:score=score1score_index=score_index[:score.size]out = score / tempv= out[:min(top_k, score.size)]idx_next = torch.multinomial(torch.Tensor(v), num_samples=1, generator=None)if voc["voc"][score_index[idx_next.item()]] == "<|sos|>":breakprompt += [voc["voc"][score_index[idx_next.item()]]]print(prompt[-1], end="", flush=True)def t_infre():model, voc = load_model_and_voc()while True:text = input("user:")gen_token(voc, model, ["<|user|>"] + list("{}".format(text)) + ["<|agent|>"], 64)print()if __name__ == '__main__':# print(pd.read_pickle("loss916"))# gen_one_voc()# gen_voc()# for i in range(17,18):# gen_pre_data_align(i, 16)# train()# gen_sft_single_data_align()# train_single()# sft 推理 一本正经的胡说八道已练成t_infre()