在使用transformers库中的模型生成时,调用generate方法并指定prefix_allowed_tokens_fn,限定模型的输出
model.generate(
input_ids=input_ids,
attention_mask=attn,
max_length=30,
prefix_allowed_tokens_fn=prefix_allowed_tokens,
num_beams=self.generate_num + testloader.dataset.max_positive,
num_return_sequences=self.generate_num + testloader.dataset.max_positive,
output_scores=True,
return_dict_in_generate=True,
)
运行后偶尔报错:
ValueError: `prefix_allowed_tokens_fn` returned an empty list for batch ID 0.This means that the constraint is unsatisfiable. Please check your implementationof `prefix_allowed_tokens_fn`
其中prefix_allowed_tokens的实现为
class Trie(object):
def __init__(self, sequences: List[List[int]] = []):
self.trie_dict = {}
self.len = 0
if sequences:
for sequence in sequences:
Trie._add_to_trie(sequence, self.trie_dict)
self.len += 1
self.append_trie = None
self.bos_token_id = None
def append(self, trie, bos_token_id):
self.append_trie = trie
self.bos_token_id = bos_token_id
def add(self, sequence: List[int]):
Trie._add_to_trie(sequence, self.trie_dict)
self.len += 1
def get(self, prefix_sequence: List[int]):
return Trie._get_from_trie(
prefix_sequence, self.trie_dict, self.append_trie, self.bos_token_id
)
@staticmethod
def load_from_dict(trie_dict):
trie = Trie()
trie.trie_dict = trie_dict
trie.len = sum(1 for _ in trie)
return trie
@staticmethod
def _add_to_trie(sequence: List[int], trie_dict: Dict):
if sequence:
if sequence[0] not in trie_dict:
trie_dict[sequence[0]] = {}
Trie._add_to_trie(sequence[1:], trie_dict[sequence[0]])
@staticmethod
def _get_from_trie(
prefix_sequence: List[int],
trie_dict: Dict,
append_trie=None,
bos_token_id: int = None,
):
if len(prefix_sequence) == 0:
output = list(trie_dict.keys())
if append_trie and bos_token_id in output:
output.remove(bos_token_id)
output += list(append_trie.trie_dict.keys())
return output
elif prefix_sequence[0] in trie_dict:
return Trie._get_from_trie(
prefix_sequence[1:],
trie_dict[prefix_sequence[0]],
append_trie,
bos_token_id,
)
else:
if append_trie:
return append_trie.get(prefix_sequence)
else:
return []
def __iter__(self):
def _traverse(prefix_sequence, trie_dict):
if trie_dict:
for next_token in trie_dict:
yield from _traverse(
prefix_sequence + [next_token], trie_dict[next_token]
)
else:
yield prefix_sequence
return _traverse([], self.trie_dict)
def __len__(self):
return self.len
def __getitem__(self, value):
return self.get(value)
def prefix_allowed_tokens_fn(candidate_trie):
def prefix_allowed_tokens(batch_id, sentence):
sentence = sentence.tolist()
trie_out = candidate_trie.get(sentence)
return trie_out
return prefix_allowed_tokens
candidate_trie = gt.Trie(
[
[0] + self.tokenizer.encode(f"{candidate}")
for candidate in candidates
]
)
prefix_allowed_tokens = prefix_allowed_tokens_fn(candidate_trie)
经过各种排查,最终确定问题实际上在于transformers的版本,经过测试transformers==4.26.0时可以正常运行。transformers==4.40及以上版本会报错