模型自动采样

张开发
2026/4/10 23:40:34 15 分钟阅读

分享文章

模型自动采样
import torch from transformers import AutoTokenizer, AutoModelForCausalLM # 加载最轻量级的模型distilgpt2 model_name distilgpt2 # 仅 82M 参数速度快资源少 tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained(model_name) # 设置模型为评估模式不训练 model.eval() prompt The weather today is inputs tokenizer(prompt, return_tensorspt) # 生成 2 条不同的回复每条最多 20 个新 token outputs model.generate( **inputs, max_new_tokens20, do_sampleTrue, # 启用采样 temperature0.8, # 控制随机性1 更随机1 更确定 top_k50, # 只从概率最高的 50 个词中采样 num_return_sequences2, # 返回 2 条不同的结果 ) for i, output in enumerate(outputs): reply tokenizer.decode(output, skip_special_tokensTrue) print(f回复 {i1}: {reply})

更多文章