使用Hugging Face Transformers计算token数时如何屏蔽超限警告？

问题 #

我需要计算字符串的token数量，方便在发送给OpenAI前判断是否超出模型最大长度，超出就拆分。目前用以下代码实现：

from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
return len(tokenizer.encode(text))

代码能正常运行，但控制台频繁弹出警告：

Token indices sequence length is longer than the specified maximum sequence length for this model (2551 > 1024). Running this sequence through the model will result in indexing errors

我知道文本长度已经超过模型最大长度，这正是我用这个方法的原因，想知道怎么忽略这个警告，或者有没有更优的实现方式，避免警告刷屏影响查看重要日志。

解决方案 #

方法一：精准屏蔽特定警告 #

用Python的warnings模块过滤掉该警告，有两种方式：

• 按警告类别屏蔽（需导入transformers的WarningType）：

import warnings
from transformers import GPT2Tokenizer, WarningType

warnings.filterwarnings("ignore", category=WarningType, message="Token indices sequence length is longer than the specified maximum sequence length for this model")

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
def count_tokens(text):
    return len(tokenizer.encode(text))

• 按消息关键词屏蔽（无需指定类别，匹配消息正则）：

import warnings
from transformers import GPT2Tokenizer

warnings.filterwarnings("ignore", message=r".*Token indices sequence length is longer than the specified maximum sequence length.*")

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
def count_tokens(text):
    return len(tokenizer.encode(text))

方法二：无警告的token计数实现 #

使用tokenizer的__call__方法配合return_length=True，或者直接用tokenize方法，这两种方式都不会触发长度警告：

• 用__call__方法：

from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
def count_tokens(text):
    return tokenizer(text, return_length=True).length[0]

• 用tokenize方法：

from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
def count_tokens(text):
    return len(tokenizer.tokenize(text))

注：tokenize和encode的计数结果完全一致，因为encode是先执行tokenize再转换为索引，不会影响token数量统计。

内容的提问来源于stack exchange，提问作者user11281688