54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
|
import torch
|
|||
|
|
|
|||
|
|
from gnnrec.config import DATA_DIR, MODEL_DIR
|
|||
|
|
from gnnrec.kgrec.data import OAGCSContrastDataset
|
|||
|
|
from gnnrec.kgrec.scibert import ContrastiveSciBERT
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Context:
|
|||
|
|
|
|||
|
|
def __init__(self, paper_embeds, scibert_model):
|
|||
|
|
"""论文召回模块上下文对象
|
|||
|
|
|
|||
|
|
:param paper_embeds: tensor(N, d) 论文标题向量
|
|||
|
|
:param scibert_model: ContrastiveSciBERT 微调后的SciBERT模型
|
|||
|
|
"""
|
|||
|
|
self.paper_embeds = paper_embeds
|
|||
|
|
self.scibert_model = scibert_model
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_context():
|
|||
|
|
paper_embeds = torch.load(DATA_DIR / 'oag/cs/paper_feat.pkl', map_location='cpu')
|
|||
|
|
scibert_model = ContrastiveSciBERT(128, 0.07)
|
|||
|
|
scibert_model.load_state_dict(torch.load(MODEL_DIR / 'scibert.pt', map_location='cpu'))
|
|||
|
|
return Context(paper_embeds, scibert_model)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def recall(ctx, query, k=1000):
|
|||
|
|
"""根据输入的查询词在oag-cs数据集召回论文
|
|||
|
|
|
|||
|
|
:param ctx: Context 上下文对象
|
|||
|
|
:param query: str 查询词
|
|||
|
|
:param k: int, optional 召回论文数量，默认为1000
|
|||
|
|
:return: List[float], List[int] Top k论文的相似度和id，按相似度降序排序
|
|||
|
|
"""
|
|||
|
|
q = ctx.scibert_model.get_embeds(query) # (1, d)
|
|||
|
|
q = q / q.norm()
|
|||
|
|
similarity = torch.mm(ctx.paper_embeds, q.t()).squeeze(dim=1) # (N,)
|
|||
|
|
score, pid = similarity.topk(k, dim=0)
|
|||
|
|
return score.tolist(), pid.tolist()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
ctx = get_context()
|
|||
|
|
paper_titles = OAGCSContrastDataset(DATA_DIR / 'oag/cs/mag_papers.txt', 'all')
|
|||
|
|
while True:
|
|||
|
|
query = input('query> ').strip()
|
|||
|
|
score, pid = recall(ctx, query, 10)
|
|||
|
|
for i in range(len(pid)):
|
|||
|
|
print('{:.4f}\t{}'.format(score[i], paper_titles[pid[i]][0]))
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|