503 lines
20 KiB
Python
503 lines
20 KiB
Python
from logging import debug
|
|
import random
|
|
import pytorch_lightning as pl
|
|
import torch
|
|
import pickle
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import json
|
|
# from transformers.utils.dummy_pt_objects import PrefixConstrainedLogitsProcessor
|
|
|
|
from .base import BaseLitModel
|
|
from transformers.optimization import get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup
|
|
|
|
from functools import partial
|
|
from .utils import rank_score, acc, LabelSmoothSoftmaxCEV1
|
|
|
|
from typing import Callable, Iterable, List
|
|
|
|
def lmap(f: Callable, x: Iterable) -> List:
|
|
"""list(map(f, x))"""
|
|
return list(map(f, x))
|
|
|
|
def multilabel_categorical_crossentropy(y_pred, y_true):
|
|
y_pred = (1 - 2 * y_true) * y_pred
|
|
y_pred_neg = y_pred - y_true * 1e12
|
|
y_pred_pos = y_pred - (1 - y_true) * 1e12
|
|
zeros = torch.zeros_like(y_pred[..., :1])
|
|
y_pred_neg = torch.cat([y_pred_neg, zeros], dim=-1)
|
|
y_pred_pos = torch.cat([y_pred_pos, zeros], dim=-1)
|
|
neg_loss = torch.logsumexp(y_pred_neg, dim=-1)
|
|
pos_loss = torch.logsumexp(y_pred_pos, dim=-1)
|
|
return (neg_loss + pos_loss).mean()
|
|
|
|
def decode(output_ids, tokenizer):
|
|
return lmap(str.strip, tokenizer.batch_decode(output_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True))
|
|
|
|
class TransformerLitModel(BaseLitModel):
|
|
def __init__(self, model, args, tokenizer=None, data_config={}):
|
|
super().__init__(model, args)
|
|
self.save_hyperparameters(args)
|
|
if args.bce:
|
|
self.loss_fn = torch.nn.BCEWithLogitsLoss()
|
|
elif args.label_smoothing != 0.0:
|
|
self.loss_fn = LabelSmoothSoftmaxCEV1(lb_smooth=args.label_smoothing)
|
|
else:
|
|
self.loss_fn = nn.CrossEntropyLoss()
|
|
self.best_acc = 0
|
|
self.first = True
|
|
|
|
self.tokenizer = tokenizer
|
|
|
|
self.__dict__.update(data_config)
|
|
# resize the word embedding layer
|
|
self.model.resize_token_embeddings(len(self.tokenizer))
|
|
self.decode = partial(decode, tokenizer=self.tokenizer)
|
|
if args.pretrain:
|
|
self._freaze_attention()
|
|
elif "ind" in args.data_dir:
|
|
# for inductive setting, use feeaze the word embedding
|
|
self._freaze_word_embedding()
|
|
|
|
|
|
def forward(self, x):
|
|
return self.model(x)
|
|
|
|
def training_step(self, batch, batch_idx): # pylint: disable=unused-argument
|
|
# embed();exit()
|
|
# print(self.optimizers().param_groups[1]['lr'])
|
|
labels = batch.pop("labels")
|
|
label = batch.pop("label")
|
|
pos = batch.pop("pos")
|
|
try:
|
|
en = batch.pop("en")
|
|
rel = batch.pop("rel")
|
|
except KeyError:
|
|
pass
|
|
input_ids = batch['input_ids']
|
|
logits = self.model(**batch, return_dict=True).logits
|
|
|
|
_, mask_idx = (input_ids == self.tokenizer.mask_token_id).nonzero(as_tuple=True)
|
|
bs = input_ids.shape[0]
|
|
mask_logits = logits[torch.arange(bs), mask_idx][:, self.entity_id_st:self.entity_id_ed]
|
|
|
|
assert mask_idx.shape[0] == bs, "only one mask in sequence!"
|
|
if self.args.bce:
|
|
loss = self.loss_fn(mask_logits, labels)
|
|
else:
|
|
loss = self.loss_fn(mask_logits, label)
|
|
|
|
if batch_idx == 0:
|
|
print('\n'.join(self.decode(batch['input_ids'][:4])))
|
|
|
|
|
|
return loss
|
|
|
|
def _eval(self, batch, batch_idx, ):
|
|
labels = batch.pop("labels")
|
|
input_ids = batch['input_ids']
|
|
# single label
|
|
label = batch.pop('label')
|
|
pos = batch.pop('pos')
|
|
my_keys = list(batch.keys())
|
|
for k in my_keys:
|
|
if k not in ["input_ids", "attention_mask", "token_type_ids"]:
|
|
batch.pop(k)
|
|
logits = self.model(**batch, return_dict=True).logits[:, :, self.entity_id_st:self.entity_id_ed]
|
|
_, mask_idx = (input_ids == self.tokenizer.mask_token_id).nonzero(as_tuple=True)
|
|
bsz = input_ids.shape[0]
|
|
logits = logits[torch.arange(bsz), mask_idx]
|
|
# get the entity ranks
|
|
# filter the entity
|
|
assert labels[0][label[0]], "correct ids must in filiter!"
|
|
labels[torch.arange(bsz), label] = 0
|
|
assert logits.shape == labels.shape
|
|
logits += labels * -100 # mask entityj
|
|
# for i in range(bsz):
|
|
# logits[i][labels]
|
|
|
|
_, outputs = torch.sort(logits, dim=1, descending=True)
|
|
_, outputs = torch.sort(outputs, dim=1)
|
|
ranks = outputs[torch.arange(bsz), label].detach().cpu() + 1
|
|
|
|
|
|
return dict(ranks = np.array(ranks))
|
|
|
|
def validation_step(self, batch, batch_idx):
|
|
result = self._eval(batch, batch_idx)
|
|
return result
|
|
|
|
def validation_epoch_end(self, outputs) -> None:
|
|
ranks = np.concatenate([_['ranks'] for _ in outputs])
|
|
total_ranks = ranks.shape[0]
|
|
|
|
if not self.args.pretrain:
|
|
l_ranks = ranks[np.array(list(np.arange(0, total_ranks, 2)))]
|
|
r_ranks = ranks[np.array(list(np.arange(0, total_ranks, 2))) + 1]
|
|
self.log("Eval/lhits10", (l_ranks<=10).mean())
|
|
self.log("Eval/rhits10", (r_ranks<=10).mean())
|
|
|
|
hits20 = (ranks<=20).mean()
|
|
hits10 = (ranks<=10).mean()
|
|
hits3 = (ranks<=3).mean()
|
|
hits1 = (ranks<=1).mean()
|
|
|
|
self.log("Eval/hits10", hits10)
|
|
self.log("Eval/hits20", hits20)
|
|
self.log("Eval/hits3", hits3)
|
|
self.log("Eval/hits1", hits1)
|
|
self.log("Eval/mean_rank", ranks.mean())
|
|
self.log("Eval/mrr", (1. / ranks).mean())
|
|
self.log("hits10", hits10, prog_bar=True)
|
|
self.log("hits1", hits1, prog_bar=True)
|
|
|
|
|
|
|
|
|
|
def test_step(self, batch, batch_idx): # pylint: disable=unused-argument
|
|
# ranks = self._eval(batch, batch_idx)
|
|
result = self._eval(batch, batch_idx)
|
|
# self.log("Test/ranks", np.mean(ranks))
|
|
|
|
return result
|
|
|
|
def test_epoch_end(self, outputs) -> None:
|
|
ranks = np.concatenate([_['ranks'] for _ in outputs])
|
|
|
|
hits20 = (ranks<=20).mean()
|
|
hits10 = (ranks<=10).mean()
|
|
hits3 = (ranks<=3).mean()
|
|
hits1 = (ranks<=1).mean()
|
|
|
|
|
|
self.log("Test/hits10", hits10)
|
|
self.log("Test/hits20", hits20)
|
|
self.log("Test/hits3", hits3)
|
|
self.log("Test/hits1", hits1)
|
|
self.log("Test/mean_rank", ranks.mean())
|
|
self.log("Test/mrr", (1. / ranks).mean())
|
|
|
|
def configure_optimizers(self):
|
|
no_decay_param = ["bias", "LayerNorm.weight"]
|
|
|
|
optimizer_group_parameters = [
|
|
{"params": [p for n, p in self.model.named_parameters() if p.requires_grad and not any(nd in n for nd in no_decay_param)], "weight_decay": self.args.weight_decay},
|
|
{"params": [p for n, p in self.model.named_parameters() if p.requires_grad and any(nd in n for nd in no_decay_param)], "weight_decay": 0}
|
|
]
|
|
|
|
optimizer = self.optimizer_class(optimizer_group_parameters, lr=self.lr, eps=1e-8)
|
|
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=self.num_training_steps * self.args.warm_up_radio, num_training_steps=self.num_training_steps)
|
|
return {
|
|
"optimizer": optimizer,
|
|
"lr_scheduler":{
|
|
'scheduler': scheduler,
|
|
'interval': 'step', # or 'epoch'
|
|
'frequency': 1,
|
|
}
|
|
}
|
|
|
|
def _freaze_attention(self):
|
|
for k, v in self.model.named_parameters():
|
|
if "word" not in k:
|
|
v.requires_grad = False
|
|
else:
|
|
print(k)
|
|
|
|
def _freaze_word_embedding(self):
|
|
for k, v in self.model.named_parameters():
|
|
if "word" in k:
|
|
print(k)
|
|
v.requires_grad = False
|
|
|
|
@staticmethod
|
|
def add_to_argparse(parser):
|
|
parser = BaseLitModel.add_to_argparse(parser)
|
|
|
|
parser.add_argument("--label_smoothing", type=float, default=0.1, help="")
|
|
parser.add_argument("--bce", type=int, default=0, help="")
|
|
return parser
|
|
|
|
|
|
|
|
import faiss
|
|
import os
|
|
class GetEntityEmbeddingLitModel(TransformerLitModel):
|
|
def __init__(self, model, args, tokenizer, data_config={}):
|
|
super().__init__(model, args, tokenizer, data_config)
|
|
|
|
self.faissid2entityid = {}
|
|
# self.index = faiss.IndexFlatL2(d) # build the index
|
|
|
|
d, measure = self.model.config.hidden_size, faiss.METRIC_L2
|
|
# param = 'HNSW64'
|
|
# self.index = faiss.index_factory(d, param, measure)
|
|
self.index = faiss.IndexFlatL2(d) # build the index
|
|
# print(self.index.is_trained) # 此时输出为True
|
|
# index.add(xb)
|
|
self.cnt_batch = 0
|
|
self.total_embedding = []
|
|
|
|
|
|
|
|
def test_step(self, batch, batch_idx):
|
|
labels = batch.pop("labels")
|
|
mask_idx = batch.pop("pos")
|
|
input_ids = batch['input_ids']
|
|
# single label
|
|
label = batch.pop('label')
|
|
# last layer
|
|
hidden_states = self.model(**batch, return_dict=True, output_hidden_states=True).hidden_states[-1]
|
|
# _, mask_idx = (input_ids == self.tokenizer.mask_token_id).nonzero(as_tuple=True)
|
|
bsz = input_ids.shape[0]
|
|
entity_embedding = hidden_states[torch.arange(bsz), mask_idx].cpu()
|
|
# use normalize or not ?
|
|
# entity_embedding = F.normalize(entity_embedding, dim=-1, p = 2)
|
|
self.total_embedding.append(entity_embedding)
|
|
# self.index.add(np.array(entity_embedding, dtype=np.float32))
|
|
for i, l in zip(range(bsz), label):
|
|
self.faissid2entityid[i+self.cnt_batch] = l.cpu()
|
|
self.cnt_batch += bsz
|
|
|
|
|
|
def test_epoch_end(self, outputs) -> None:
|
|
self.total_embedding = np.concatenate(self.total_embedding, axis=0)
|
|
# self.index.train(self.total_embedding)
|
|
print(faiss.MatrixStats(self.total_embedding).comments)
|
|
self.index.add(self.total_embedding)
|
|
faiss.write_index(self.index, os.path.join(self.args.data_dir, "faiss_dump.index"))
|
|
with open(os.path.join(self.args.data_dir, "faissid2entityid.pkl") ,'wb') as file:
|
|
pickle.dump(self.faissid2entityid, file)
|
|
|
|
with open(os.path.join(self.args.data_dir, "total_embedding.pkl") ,'wb') as file:
|
|
pickle.dump(self.total_embedding, file)
|
|
# print(f"number of entity embedding : {len(self.faissid2entityid)}")
|
|
|
|
@staticmethod
|
|
def add_to_argparse(parser):
|
|
parser = TransformerLitModel.add_to_argparse(parser)
|
|
parser.add_argument("--faiss_init", type=int, default=1, help="get the embedding and save it the file.")
|
|
return parser
|
|
|
|
class UseEntityEmbeddingLitModel(TransformerLitModel):
|
|
def __init__(self, model, args, tokenizer, data_config={}):
|
|
super().__init__(model, args, tokenizer, data_config)
|
|
|
|
self.faissid2entityid = pickle.load(open(os.path.join(self.args.data_dir, "faissid2entityid.pkl") ,'rb'))
|
|
self.index = faiss.read_index(os.path.join(self.args.data_dir, "faiss_dump.index"))
|
|
|
|
|
|
self.dis2logits = distance2logits_2
|
|
|
|
def _eval(self, batch, batch_idx, ):
|
|
labels = batch.pop("labels")
|
|
pos = batch.pop("pos")
|
|
input_ids = batch['input_ids']
|
|
# single label
|
|
label = batch.pop('label')
|
|
|
|
hidden_states = self.model(**batch, return_dict=True, output_hidden_states=True).hidden_states[-1]
|
|
_, mask_idx = (input_ids == self.tokenizer.mask_token_id).nonzero(as_tuple=True)
|
|
bsz = input_ids.shape[0]
|
|
mask_embedding = np.array(hidden_states[torch.arange(bsz), mask_idx].cpu(), dtype=np.float32)
|
|
topk = 200
|
|
D, I = self.index.search(mask_embedding, topk)
|
|
labels[torch.arange(bsz), label] = 0
|
|
|
|
# logits = torch.zeros_like(labels)
|
|
# D = torch.softmax(torch.exp(-1. * torch.tensor(D)), dim=-1)
|
|
# for i in range(bsz):
|
|
# for j in range(topk):
|
|
# logits[i][self.faissid2entityid[I[i][j]]] += D[i][j]
|
|
# # get the entity ranks
|
|
# # filter the entity
|
|
# assert labels[0][label[0]], "correct ids must in filiter!"
|
|
# labels[torch.arange(bsz), label] = 0
|
|
# assert logits.shape == labels.shape
|
|
# logits += labels * -100 # mask entityj
|
|
|
|
# _, outputs = torch.sort(logits, dim=1, descending=True)
|
|
# _, outputs = torch.sort(outputs, dim=1)
|
|
# ranks = outputs[torch.arange(bsz), label].detach().cpu() + 1
|
|
|
|
|
|
entity_logits = torch.full(labels.shape, -100.).to(self.device)
|
|
D = self.dis2logits(D)
|
|
for i in range(bsz):
|
|
for j in range(topk):
|
|
# filter entity in labels
|
|
if I[i][j] not in self.faissid2entityid:
|
|
print(I[i][j])
|
|
break
|
|
# assert I[i][j] in self.faissid2entityid, print(I[i][j])
|
|
if labels[i][self.faissid2entityid[I[i][j]]]: continue
|
|
if entity_logits[i][self.faissid2entityid[I[i][j]]] == -100.:
|
|
entity_logits[i][self.faissid2entityid[I[i][j]]] = D[i][j]
|
|
# no added together
|
|
# else:
|
|
# entity_logits[i][self.faissid2entityid[I[i][j]]] += D[i][j]
|
|
# get the entity ranks
|
|
# filter the entity
|
|
|
|
assert entity_logits.shape == labels.shape
|
|
|
|
_, outputs = torch.sort(entity_logits, dim=1, descending=True)
|
|
_, outputs = torch.sort(outputs, dim=1)
|
|
ranks = outputs[torch.arange(bsz), label].detach().cpu() + 1
|
|
|
|
|
|
return dict(ranks = np.array(ranks))
|
|
|
|
|
|
@staticmethod
|
|
def add_to_argparse(parser):
|
|
parser = TransformerLitModel.add_to_argparse(parser)
|
|
parser.add_argument("--faiss_init", type=int, default=0, help="get the embedding and save it the file.")
|
|
parser.add_argument("--faiss_use", type=int, default=1, help="get the embedding and save it the file.")
|
|
return parser
|
|
|
|
|
|
class CombineEntityEmbeddingLitModel(UseEntityEmbeddingLitModel):
|
|
def __init__(self, model, args, tokenizer, data_config={}):
|
|
super().__init__(model, args, tokenizer, data_config=data_config)
|
|
self.dis2logits = distance2logits_2
|
|
self.id2entity = {}
|
|
with open("./dataset/FB15k-237/entity2textlong.txt", 'r') as file:
|
|
cnt = 0
|
|
for line in file.readlines():
|
|
e, d = line.strip().split("\t")
|
|
self.id2entity[cnt] = e
|
|
cnt += 1
|
|
self.id2entity_t = {}
|
|
with open("./dataset/FB15k-237/entity2text.txt", 'r') as file:
|
|
for line in file.readlines():
|
|
e, d = line.strip().split("\t")
|
|
self.id2entity_t[e] = d
|
|
for k, v in self.id2entity.items():
|
|
self.id2entity[k] = self.id2entity_t[v]
|
|
def _eval(self, batch, batch_idx, ):
|
|
labels = batch.pop("labels")
|
|
input_ids = batch['input_ids']
|
|
# single label
|
|
label = batch.pop('label')
|
|
pos = batch.pop("pos")
|
|
|
|
result = self.model(**batch, return_dict=True, output_hidden_states=True)
|
|
hidden_states = result.hidden_states[-1]
|
|
_, mask_idx = (input_ids == self.tokenizer.mask_token_id).nonzero(as_tuple=True)
|
|
bsz = input_ids.shape[0]
|
|
mask_embedding = np.array(hidden_states[torch.arange(bsz), mask_idx].cpu(), dtype=np.float32)
|
|
# mask_embedding = np.array(hidden_states[torch.arange(bsz), mask_idx].cpu(), dtype=np.float32)
|
|
topk = self.args.knn_topk
|
|
D, I = self.index.search(mask_embedding, topk)
|
|
D = torch.from_numpy(D).to(self.device)
|
|
assert labels[0][label[0]], "correct ids must in filiter!"
|
|
labels[torch.arange(bsz), label] = 0
|
|
|
|
|
|
mask_logits = result.logits[:, :, self.entity_id_st:self.entity_id_ed]
|
|
mask_logits = mask_logits[torch.arange(bsz), mask_idx]
|
|
entity_logits = torch.full(labels.shape, 1000.).to(self.device)
|
|
# D = self.dis2logits(D)
|
|
for i in range(bsz):
|
|
for j in range(topk):
|
|
# filter entity in labels
|
|
if labels[i][self.faissid2entityid[I[i][j]]]: continue
|
|
if entity_logits[i][self.faissid2entityid[I[i][j]]] == 1000.:
|
|
entity_logits[i][self.faissid2entityid[I[i][j]]] = D[i][j]
|
|
# else:
|
|
# entity_logits[i][self.faissid2entityid[I[i][j]]] += D[i][j]
|
|
entity_logits = self.dis2logits(entity_logits)
|
|
# get the entity ranks
|
|
# filter the entity
|
|
assert entity_logits.shape == labels.shape
|
|
assert mask_logits.shape == labels.shape
|
|
# entity_logits = torch.softmax(entity_logits + labels * -100, dim=-1) # mask entityj
|
|
entity_logits = entity_logits + labels* -100.
|
|
mask_logits = torch.softmax(mask_logits + labels* -100, dim=-1)
|
|
# logits = mask_logits
|
|
logits = combine_knn_and_vocab_probs(entity_logits, mask_logits, self.args.knn_lambda)
|
|
# logits = entity_logits + mask_logits
|
|
|
|
|
|
knn_topk_logits, knn_topk_id = entity_logits.topk(20)
|
|
mask_topk_logits, mask_topk_id = mask_logits.topk(20)
|
|
union_topk = []
|
|
for i in range(bsz):
|
|
num_same = len(list(set(knn_topk_id[i].cpu().tolist()) & set(mask_topk_id[i].cpu().tolist())))
|
|
union_topk.append(num_same/ 20.)
|
|
|
|
knn_topk_id = knn_topk_id.to("cpu")
|
|
mask_topk_id = mask_topk_id.to("cpu")
|
|
mask_topk_logits = mask_topk_logits.to("cpu")
|
|
knn_topk_logits = knn_topk_logits.to("cpu")
|
|
label = label.to("cpu")
|
|
|
|
|
|
|
|
for t in range(bsz):
|
|
if knn_topk_id[t][0] == label[t] and knn_topk_logits[t][0] > mask_topk_logits[t][0] and mask_topk_logits[t][0] <= 0.4:
|
|
print(knn_topk_logits[t], knn_topk_id[t])
|
|
print(lmap(lambda x: self.id2entity[x.item()], knn_topk_id[t]))
|
|
print(mask_topk_logits[t], mask_topk_id[t])
|
|
print(lmap(lambda x: self.id2entity[x.item()], mask_topk_id[t]))
|
|
print(label[t])
|
|
print()
|
|
|
|
_, outputs = torch.sort(logits, dim=1, descending=True)
|
|
_, outputs = torch.sort(outputs, dim=1)
|
|
ranks = outputs[torch.arange(bsz), label].detach().cpu() + 1
|
|
|
|
|
|
return dict(ranks = np.array(ranks), knn_topk_id=knn_topk_id, knn_topk_logits=knn_topk_logits,
|
|
mask_topk_id=mask_topk_id, mask_topk_logits=mask_topk_logits, num_same = np.array(union_topk))
|
|
|
|
def test_epoch_end(self, outputs) -> None:
|
|
|
|
ranks = np.concatenate([_['ranks'] for _ in outputs])
|
|
num_same = np.concatenate([_['num_same'] for _ in outputs])
|
|
results_keys = list(outputs[0].keys())
|
|
results = {}
|
|
# for k in results_keys:
|
|
# results.
|
|
|
|
self.log("Test/num_same", num_same.mean())
|
|
|
|
hits20 = (ranks<=20).mean()
|
|
hits10 = (ranks<=10).mean()
|
|
hits3 = (ranks<=3).mean()
|
|
hits1 = (ranks<=1).mean()
|
|
|
|
|
|
self.log("Test/hits10", hits10)
|
|
self.log("Test/hits20", hits20)
|
|
self.log("Test/hits3", hits3)
|
|
self.log("Test/hits1", hits1)
|
|
self.log("Test/mean_rank", ranks.mean())
|
|
self.log("Test/mrr", (1. / ranks).mean())
|
|
|
|
def add_to_argparse(parser):
|
|
parser = TransformerLitModel.add_to_argparse(parser)
|
|
parser.add_argument("--knn_lambda", type=float, default=0.5, help="lambda * knn + (1-lambda) * mask logits , lambda of knn logits and mask logits.")
|
|
parser.add_argument("--knn_topk", type=int, default=100, help="")
|
|
|
|
return parser
|
|
|
|
def combine_knn_and_vocab_probs(knn_p, vocab_p, coeff=0.5):
|
|
combine_probs = torch.stack([vocab_p, knn_p], dim=0)
|
|
coeffs = torch.ones_like(combine_probs)
|
|
coeffs[0] = np.log(1 - coeff)
|
|
coeffs[1] = np.log(coeff)
|
|
curr_prob = torch.logsumexp(combine_probs + coeffs, dim=0)
|
|
|
|
return curr_prob
|
|
|
|
def distance2logits(D):
|
|
return torch.softmax( -1. * torch.tensor(D) / 30., dim=-1)
|
|
|
|
def distance2logits_2(D, n=10):
|
|
if not isinstance(D, torch.Tensor):
|
|
D = torch.tensor(D)
|
|
if torch.sum(D) != 0.0:
|
|
distances = torch.exp(-D/n) / torch.sum(torch.exp(-D/n), dim=-1, keepdim=True)
|
|
return distances |