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上運行代碼。