import numpy as np
import torch
from torch.utils.data import Dataset
classDataReader:
NEGATIVE_TABLE_SIZE =1e8def__init__(self,dataset,min_count,care_type):
self.negatives =[]
self.discards =[]
self.negpos =0
self.care_type = care_type
self.word2id =dict()
self.id2word =dict()
self.sentences_count =0
self.token_count =0
self.word_frequency =dict()
self.inputFileName = dataset.fn
self.read_words(min_count)
self.initTableNegatives()
self.initTableDiscards()defread_words(self,min_count):
word_frequency =dict()for line inopen(self.inputFileName,encoding ='ISO-8859-1'):
line = line.split()iflen(line)>1:
self.sentences_count +=1for word in line:iflen(word)>0:
self.token_count +=1
word_frequency[word]= word_frequency.get(word,0)+1if self.token_count %1000000==0:print("Read "+str(int(self.token_count /1000000))+"M words.")
wid =0for w,c in word_frequency.items():if c < min_count:continue
self.word2id[w]= wid
self.id2word[wid]= w
self.word_frequency[wid]= c
wid +=1
self.word_count =len(self.word2id)print('Total embeddings:'+str(len(self.word2id)))definitTableNegatives(self):
pow_frequency = np.array(list(self.word_frequency.values()))**0.75
words_pow =sum(pow_frequency)
ratio = pow_frequency / words_pow
count = np.round(ratio * DataReader.NEGATIVE_TABLE_SIZE)for wid,c inenumerate(count):
self.negatives +=[wid]*int(c)
self.negatives = np.array(self.negatives)
np.random.shuffle(self.negatives)
self.sampling_prob = ratio
definitTableDiscards(self):
t =0.0001
f = np.array(list(self.word_frequency.values()))/ self.token_count
self.discards = np.sqrt(t / f)+(t / f)defgetNegatives(self,target,size):if self.care_type ==0:
response = self.negatives[self.negpos:self.negpos + size]
self.negpos =(self.negpos + size)%len(self.negatives)iflen(response)!=size:return np.concatenate((response,self.negatives[:self.negpos]))return response
5.dataset定义
classMetapath2vecDataset(Dataset):def__init__(self,data,window_size):
self.data = data
self.window_size = window_size
self.input_file =open(data.inputFileName,encoding='ISO-8859-1')def__len__(self):return self.data.sentences_count
def__getitem__(self,idx):whileTrue:
line = self.input_file.readline()ifnot line:
self.input_file.seek(0,0)
line = self.input_file.readline()iflen(line)>1:
words = line.split()iflen(words)>1:
word_ids =[self.data.word2id[w]for w in words if
w in self.data.word2id and np.random.rand()< self.data.discards[self.data.word2id[w]]]
pair_catch =[]for i,u inenumerate(word_ids):for j,v inenumerate(
word_ids[max(i - self.window_size,0):i + self.window_size]):assert u < self.data.word_count
assert v < self.data.word_count
if i == j:continue
pair_catch.append((u,v,self.data.getNegatives(v,5)))return pair_catch
@staticmethoddefcollate(batches):
all_u =[u for batch in batches for u,_,_ in batch iflen(batch)>0]
all_v =[v for batch in batches for _,v,_ in batch iflen(batch)>0]
all_neg_v =[neg_v for batch in batches for _,_,neg_v in batch iflen(batch)>0]assertlen(all_u)==len(all_v)==len(all_neg_v)# print(1)return torch.LongTensor(all_u),torch.LongTensor(all_v),torch.LongTensor(all_neg_v)'''
dataset = Metapath2vecDataset(data,3)
dataloader = DataLoader(dataset,shuffle=True,batch_size = args.batch_size,collate_fn=dataset.collate)
for v in dataloader:
print(v[0].shape)
print(v[1].shape)
print(v[2].shape)
break
'''