31 lines
990 B
Python
31 lines
990 B
Python
from torch.utils.data import Dataset
|
||
|
||
from gnnrec.kgrec.utils import iter_json
|
||
|
||
|
||
class OAGCSContrastDataset(Dataset):
|
||
SPLIT_YEAR = 2016
|
||
|
||
def __init__(self, raw_file, split='train'):
|
||
"""oag-cs论文标题-关键词对比学习数据集
|
||
|
||
由于原始数据不包含关键词，因此使用研究领域（fos字段）作为关键词
|
||
|
||
:param raw_file: str 原始论文数据文件
|
||
:param split: str "train", "valid", "all"
|
||
"""
|
||
self.titles = []
|
||
self.keywords = []
|
||
for p in iter_json(raw_file):
|
||
if split == 'train' and p['year'] <= self.SPLIT_YEAR \
|
||
or split == 'valid' and p['year'] > self.SPLIT_YEAR \
|
||
or split == 'all':
|
||
self.titles.append(p['title'])
|
||
self.keywords.append('; '.join(p['fos']))
|
||
|
||
def __getitem__(self, item):
|
||
return self.titles[item], self.keywords[item]
|
||
|
||
def __len__(self):
|
||
return len(self.titles)
|