点击查看代码
from transformers import AutoModelForCausalLM, AutoTokenizer
import torchmodel_path = "someplace"
device = "cuda:0"tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path,device_map=device,torch_dtype=torch.bfloat16 #表示使用16位浮点数据类型
)system_prompt = "there is system_prompt"
prompt = "there is prompt"inputs = tokenizer.apply_chat_template([{"role": "system", "content": system_prompt},{"role": "user", "content": prompt}],tokenizer=True, # 是否进行标记化add_generation_prompt=True, # 是否添加生成提示return_tensors="pt", # 返回格式("pt"/"tf"/"np")chat_template=None, # 自定义聊天模板padding=False, # 是否对批次进行填充truncation=False, # 是否对超长序列截断max_length=None
).to(device)'''
完整的apply_chat_template参数为：def apply_chat_template(self,conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]]], # 对话历史记录，包含角色和内容tools: Optional[List[Union[Dict, Callable]]] = None, # 可选参数，提供模型可调用的工具（函数）列表，用于支持函数调用功能，如搜索、计算等，可以是JSON模式或可调用函数documents: Optional[List[Dict[str, str]]] = None, # 可选参数，RAG（检索增强生成）所需的文档列表，通常包含"title"和"text"键的字典列表chat_template: Optional[str] = None, add_generation_prompt: bool = False,continue_final_message: bool = False, # 设置为True时，会格式化最后一条消息为开放式（无结束标记），允许模型继续该消息，适合预填充回复的场景，与add_generation_prompt互斥输出处理参数tokenize: bool = True, padding: Union[bool, str, PaddingStrategy] = False, # 控制填充策略（不填充、填充到最长序列、填充到指定长度等）truncation: bool = False, # 是否对序列进行截断处理max_length: Optional[int] = None, # 填充或截断的最大长度return_tensors: Optional[Union[str, TensorType]] = None,return_dict: bool = False, # 是否返回带命名输出的字典return_assistant_tokens_mask: bool = False, # 是否返回标识助手生成token的掩码tokenizer_kwargs: Optional[Dict[str, Any]] = None, # 传递给tokenizer的额外参数**kwargs,)
'''generate_ids = model.generate(inputs,max_new_tokens=1000,do_sample=False, # 关闭随机采样temperature=0 # 控制生成的随机性
)'''
还有这些参数可以设置
generate_ids = model.generate(inputs, min_length=50, # 生成至少 50 个 tokenmin_new_tokens=10, # 生成至少 10 个新 tokenrepetition_penalty=1.2, # 重复惩罚no_repeat_ngram_size=2, # 对应size的ngram只能出现1次temperature=0.7, # 控制随机性top_k=50, # 只考虑前 50 个最高概率的 tokentop_p=0.9, # 只选取累计概率 90% 的 token 进行采样do_sample=True # 启用随机采样
)'''response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
