# -*- coding:utf-8 -*-
import os
from typing import Union, List
from ptune_config import *
pc = ProjectConfig()
class Verbalizer(object):
"""
Verbalizer类，用于将一个Label对应到其子Label的映射。
"""
def __init__(self, verbalizer_file: str, tokenizer, max_label_len: int):
"""
Args:
verbalizer_file (str): verbalizer文件存放地址。
tokenizer: 用于文本和id之间的转换。
max_label_len (int): 标签长度，若大于则截断，若小于则补齐
"""
self.tokenizer = tokenizer
self.label_dict = self.load_label_dict(verbalizer_file)
self.max_label_len = max_label_len
def load_label_dict(self, verbalizer_file: str) -> dict:
"""
读取本地文件，构建verbalizer字典。
Args:
verbalizer_file (str): verbalizer文件存放地址。
Returns:
dict -> {
'体育': ['篮球', '足球','网球', '排球', ...],
'酒店': ['宾馆', '旅馆', '旅店', '酒店', ...],
...
}
"""
label_dict = {}
with open(verbalizer_file, 'r', encoding='utf8') as f:
for line in f.readlines():
label, sub_labels = line.strip().split('\t')
label_dict[label] = list(set(sub_labels.split(',')))
return label_dict
def find_sub_labels(self, label: Union[list, str]) -> dict:
"""
通过标签找到所有的子标签。
Args:
label (Union[list, str]): 标签, 文本型 或 id_list, e.g. -> '体育' or [860, 5509]
Returns:
dict -> {
'sub_labels': ['笔记本', '电脑'],
'token_ids': [[5011, 6381, 3315], [4510, 5554]]
}
"""
if type(label) == list: # 如果传入为id_list, 则通过tokenizer转回来
while self.tokenizer.pad_token_id in label:
label.remove(self.tokenizer.pad_token_id)
label = ''.join(self.tokenizer.convert_ids_to_tokens(label))
# print(f'label-->{label}')
if label not in self.label_dict:
raise ValueError(f'Lable Error: "{label}" not in label_dict {list(self.label_dict)}.')
sub_labels = self.label_dict[label]
ret = {'sub_labels': sub_labels}
token_ids = [_id[1:-1] for _id in self.tokenizer(sub_labels)['input_ids']]
# print(f'token_ids-->{token_ids}')
for i in range(len(token_ids)):
token_ids[i] = token_ids[i][:self.max_label_len] # 对标签进行截断与补齐
if len(token_ids[i]) < self.max_label_len:
token_ids[i] = token_ids[i] + [self.tokenizer.pad_token_id] * (self.max_label_len - len(token_ids[i]))
ret['token_ids'] = token_ids
return ret
def batch_find_sub_labels(self, label: List[Union[list, str]]) -> list:
"""
批量找到子标签。
Args:
label (List[list, str]): 标签列表, [[4510, 5554], [860, 5509]] or ['体育', '电脑']
Returns:
list -> [
{
'sub_labels': ['笔记本', '电脑'],
'token_ids': [[5011, 6381, 3315], [4510, 5554]]
},
...
]
"""
return [self.find_sub_labels(l) for l in label]
def get_common_sub_str(self, str1: str, str2: str):
"""
寻找最大公共子串。
str1:abcd
str2:abadbcdba
"""
lstr1, lstr2 = len(str1), len(str2)
# 生成0矩阵，为方便后续计算，比字符串长度多了一列
record = [[0 for i in range(lstr2 + 1)] for j in range(lstr1 + 1)]
p = 0 # 最长匹配对应在str1中的最后一位
maxNum = 0 # 最长匹配长度
for i in range(lstr1):
for j in range(lstr2):
if str1[i] == str2[j]:
record[i+1][j+1] = record[i][j] + 1
if record[i+1][j+1] > maxNum:
maxNum = record[i+1][j+1]
p = i + 1
return str1[p-maxNum:p], maxNum
def hard_mapping(self, sub_label: str) -> str:
"""
强匹配函数，当模型生成的子label不存在时，通过最大公共子串找到重合度最高的主label。
Args:
sub_label (str): 子label。
Returns:
str: 主label。
"""
label, max_overlap_str = '', 0
# print(self.label_dict.items())
for main_label, sub_labels in self.label_dict.items():
overlap_num = 0
for s_label in sub_labels: # 求所有子label与当前推理label之间的最长公共子串长度
# print(f'self.get_common_sub_str(sub_label, s_label)-->{self.get_common_sub_str(sub_label, s_label)}')
overlap_num += self.get_common_sub_str(sub_label, s_label)[1]
if overlap_num >= max_overlap_str:
max_overlap_str = overlap_num
label = main_label
return label
def find_main_label(self, sub_label: List[Union[list, str]], hard_mapping=True) -> dict:
"""
通过子标签找到父标签。
Args:
sub_label (List[Union[list, str]]): 子标签, 文本型 或 id_list, e.g. -> '苹果' or [5741, 3362]
hard_mapping (bool): 当生成的词语不存在时，是否一定要匹配到一个最相似的label。
Returns:
dict -> {
'label': '水果',
'token_ids': [3717, 3362]
}
"""
if type(sub_label) == list: # 如果传入为id_list, 则通过tokenizer转回来
pad_token_id = self.tokenizer.pad_token_id
while pad_token_id in sub_label: # 移除[PAD]token
sub_label.remove(pad_token_id)
sub_label = ''.join(self.tokenizer.convert_ids_to_tokens(sub_label))
# print(sub_label)
main_label = '无'
for label, s_labels in self.label_dict.items():
if sub_label in s_labels:
main_label = label
break
if main_label == '无' and hard_mapping:
main_label = self.hard_mapping(sub_label)
# print(main_label)
ret = {
'label': main_label,
'token_ids': self.tokenizer(main_label)['input_ids'][1:-1]
}
return ret
def batch_find_main_label(self, sub_label: List[Union[list, str]], hard_mapping=True) -> list:
"""
批量通过子标签找父标签。
Args:
sub_label (List[Union[list, str]]): 子标签列表, ['苹果', ...] or [[5741, 3362], ...]
Returns:
list: [
{
'label': '水果',
'token_ids': [3717, 3362]
},
...
]
"""
return [self.find_main_label(l, hard_mapping) for l in sub_label]
if __name__ == '__main__':
from rich import print
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)
verbalizer = Verbalizer(
verbalizer_file=pc.verbalizer,
tokenizer=tokenizer,
max_label_len=2
)
print(verbalizer.label_dict)
# label = ['电脑', '衣服']
label = [[4510, 5554], [6132, 3302]]
ret = verbalizer.batch_find_sub_labels(label)
print(ret)