LLM调用记录

最近做了点与LLM相关的Project，于是乎记录下一点点基础代码，欢迎读者/自己随时补充

点击查看代码

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_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}
    ],
    tokenize=True,              # 是否对输出转化为token ID
    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 个 token
    min_new_tokens=10,           # 生成至少 10 个新 token
    repetition_penalty=1.2,      # 重复惩罚
    no_repeat_ngram_size=2,      # 对应size的ngram只能出现1次
    temperature=0.7,             # 控制随机性
    top_k=50,                    # 只考虑前 50 个最高概率的 token
    top_p=0.9,                   # 只选取累计概率 90% 的 token 进行采样
    do_sample=True               # 启用随机采样
)

'''

response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

if "<|Assistant|>" in response:
  response = response.split("<|Assistant|>", 1)[1].strip()