| 535 | |
| 536 | |
| 537 | class Title_W2V_Dataset(Dataset): |
| 538 | def __init__(self, path_vid, wv_from_text): |
| 539 | self.data_complete = pd.read_json('./data/data.json',orient='records',dtype=False,lines=True) |
| 540 | |
| 541 | self.vid = [] |
| 542 | with open('./data/vids/'+path_vid, "r") as fr: |
| 543 | for line in fr.readlines(): |
| 544 | self.vid.append(line.strip()) |
| 545 | self.data = self.data_complete[self.data_complete.video_id.isin(self.vid)] |
| 546 | self.data['video_id'] = self.data['video_id'].astype('category') |
| 547 | self.data['video_id'].cat.set_categories(self.vid, inplace=True) |
| 548 | self.data.sort_values('video_id', ascending=True, inplace=True) |
| 549 | self.data.reset_index(inplace=True) |
| 550 | |
| 551 | self.wv_from_text = wv_from_text |
| 552 | |
| 553 | def __len__(self): |
| 554 | return self.data.shape[0] |
| 555 | |
| 556 | def __getitem__(self, idx): |
| 557 | item = self.data.iloc[idx] |
| 558 | |
| 559 | label = 1 if item['annotation']=='假' else 0 |
| 560 | label = torch.tensor(label) |
| 561 | |
| 562 | text = item['description']+' '+item['ocr'] |
| 563 | title_w2v = [] |
| 564 | for word in jieba.cut(text, cut_all=False): |
| 565 | if self.wv_from_text.__contains__(word): |
| 566 | try: |
| 567 | title_w2v.append(self.wv_from_text[word]) |
| 568 | except: |
| 569 | continue |
| 570 | |
| 571 | title_w2v = torch.FloatTensor(title_w2v) |
| 572 | |
| 573 | return { |
| 574 | 'label': label, |
| 575 | 'title_w2v': title_w2v, |
| 576 | } |
| 577 | |
| 578 | |
| 579 | class CommentsDataset(Dataset): |