Kgt5 Wikikg90mv2

🚀 T5小模型（基于WikiKG90Mv2数据集）

这是一个在WikiKG90Mv2数据集上从头开始训练的t5-small模型。该模型主要用于尾实体预测任务，即给定主题实体和关系，预测对象实体。

🚀 快速开始

本模型是在WikiKG90Mv2数据集上进行训练的，训练约1.5个轮次，使用4块1080Ti GPU，每个轮次的训练时间约为5.5天。要评估该模型，可对每个输入 (s, r) 对从解码器采样300次，然后移除无法映射回有效实体的预测结果，并根据对数概率对预测结果进行排序，最后进行过滤。该模型在验证集上的MRR为0.22。

✨ 主要特性

• 训练数据：使用WikiKG90Mv2数据集进行训练。
• 训练任务：专注于尾实体预测任务。
• 评估方法：对每个输入采样300次，过滤无效预测并排序。

📦 安装指南

你可以使用以下代码在ipython笔记本中评估预训练模型：

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
tokenizer = AutoTokenizer.from_pretrained("apoorvumang/kgt5-wikikg90mv2")
model = AutoModelForSeq2SeqLM.from_pretrained("apoorvumang/kgt5-wikikg90mv2")

💻 使用示例

基础用法

# 一个验证集中模型预测正确的示例
# 你可以在这里尝试自己的示例。你的贵族头衔是什么？
input = "Sophie Valdemarsdottir| noble title"
out = topkSample(input, model, tokenizer, num_samples=5)
out

高级用法

以下是一些辅助函数的代码，用于获取序列分数、采样和贪婪预测：

import torch

def getScores(ids, scores, pad_token_id):
    """get sequence scores from model.generate output"""
    scores = torch.stack(scores, dim=1)
    log_probs = torch.log_softmax(scores, dim=2)
    # remove start token
    ids = ids[:,1:]
    # gather needed probs
    x = ids.unsqueeze(-1).expand(log_probs.shape)
    needed_logits = torch.gather(log_probs, 2, x)
    final_logits = needed_logits[:, :, 0]
    padded_mask = (ids == pad_token_id)
    final_logits[padded_mask] = 0
    final_scores = final_logits.sum(dim=-1)
    return final_scores.cpu().detach().numpy()

def topkSample(input, model, tokenizer, 
                num_samples=5,
                num_beams=1,
                max_output_length=30):
    tokenized = tokenizer(input, return_tensors="pt")
    out = model.generate(**tokenized,
                        do_sample=True,
                        num_return_sequences = num_samples,
                        num_beams = num_beams,
                        eos_token_id = tokenizer.eos_token_id,
                        pad_token_id = tokenizer.pad_token_id,
                        output_scores = True,
                        return_dict_in_generate=True,
                        max_length=max_output_length,)
    out_tokens = out.sequences
    out_str = tokenizer.batch_decode(out_tokens, skip_special_tokens=True)
    out_scores = getScores(out_tokens, out.scores, tokenizer.pad_token_id)
    
    pair_list = [(x[0], x[1]) for x in zip(out_str, out_scores)]
    sorted_pair_list = sorted(pair_list, key=lambda x:x[1], reverse=True)
    return sorted_pair_list

def greedyPredict(input, model, tokenizer):
    input_ids = tokenizer([input], return_tensors="pt").input_ids
    out_tokens = model.generate(input_ids)
    out_str = tokenizer.batch_decode(out_tokens, skip_special_tokens=True)
    return out_str[0]

🔧 技术细节

• 实体和关系表示：使用原始文本标题和描述来获取实体和关系的文本表示。这些原始文本来自ogb数据集本身（dataset/wikikg90m-v2/mapping/entity.csv和relation.csv）。实体表示设置为标题，如果两个实体具有相同的标题，则使用描述进行消歧。如果仍然无法消歧，则使用wikidata ID（例如Q123456）。
• 评估过程：对每个输入 (s, r) 对从解码器采样300次，移除无法映射回有效实体的预测结果，根据对数概率对预测结果进行排序，然后进行过滤。

📚 详细文档

你可以从以下链接获取更多详细信息：

• 方法详情：https://github.com/apoorvumang/kgt5/
• 完整排行榜：https://ogb.stanford.edu/docs/lsc/leaderboards/#wikikg90mv2

📄 许可证

本项目采用MIT许可证。

其他补充说明

你可以进一步加载实体别名列表，然后过滤出有效的实体预测，再创建从别名到整数ID的反向映射，以获得所需格式的最终预测结果。不过，将这些别名作为字典加载到内存中需要大量的RAM，并且你需要下载别名文件（可从https://storage.googleapis.com/kgt5-wikikg90mv2/ent_alias_list.pickle 下载，关系文件：https://storage.googleapis.com/kgt5-wikikg90mv2/rel_alias_list.pickle）。

以下是一些用于下载验证集文件和评估模型的代码示例：

# download valid.txt. you can also try same url with test.txt. however test does not contain the correct tails
!wget https://storage.googleapis.com/kgt5-wikikg90mv2/valid.txt

fname = 'valid.txt'
valid_lines = []
f = open(fname)
for line in f:
    valid_lines.append(line.rstrip())
f.close()
print(valid_lines[0])

from tqdm.auto import tqdm
# try unfiltered hits@k. this is approximation since model can sample same seq multiple times
# you should run this on gpu if you want to evaluate on all points with 300 samples each
k = 1
count_at_k = 0
max_predictions = k
max_points = 1000
for line in tqdm(valid_lines[:max_points]):
    input, target = line.split('\t')
    model_output = topkSample(input, model, tokenizer, num_samples=max_predictions)
    prediction_strings = [x[0] for x in model_output]
    if target in prediction_strings:
        count_at_k += 1
print('Hits at {0} unfiltered: {1}'.format(k, count_at_k/max_points))

⚠️ 重要提示

完整的实体映射到ID、过滤等过程为了简化未包含在此处，但如有需要可按需提供。如有疑问或需要详细信息，请联系Apoorv (apoorvumang@gmail.com)。

💡 使用建议

如果你想对所有点进行评估，每个点采样300次，建议在GPU上运行代码。