| 577 | |
| 578 | |
| 579 | class CommentsDataset(Dataset): |
| 580 | |
| 581 | def __init__(self, path_vid): |
| 582 | self.data_complete = pd.read_json('./data/data.json',orient='records',dtype=False,lines=True) |
| 583 | |
| 584 | self.vid = [] |
| 585 | with open('./data/vids/'+path_vid, "r") as fr: |
| 586 | for line in fr.readlines(): |
| 587 | self.vid.append(line.strip()) |
| 588 | self.data = self.data_complete[self.data_complete.video_id.isin(self.vid)] |
| 589 | self.data['video_id'] = self.data['video_id'].astype('category') |
| 590 | self.data['video_id'].cat.set_categories(self.vid, inplace=True) |
| 591 | self.data.sort_values('video_id', ascending=True, inplace=True) |
| 592 | self.data.reset_index(inplace=True) |
| 593 | |
| 594 | self.tokenizer = BertTokenizer.from_pretrained('bert-base-chinese') |
| 595 | |
| 596 | hascomments = self.data['comments'].apply(lambda x:len(x)>0) |
| 597 | self.data = self.data[hascomments] |
| 598 | print (self.data.shape) |
| 599 | |
| 600 | def __len__(self): |
| 601 | return self.data.shape[0] |
| 602 | |
| 603 | def __getitem__(self, idx): |
| 604 | item = self.data.iloc[idx] |
| 605 | vid = item['video_id'] |
| 606 | |
| 607 | label = 1 if item['annotation']=='假' else 0 |
| 608 | label = torch.tensor(label) |
| 609 | |
| 610 | comments_inputid = [] |
| 611 | comments_mask = [] |
| 612 | for comment in item['comments']: |
| 613 | comment_tokens = self.tokenizer(comment, max_length=250, padding='max_length', truncation=True) |
| 614 | comments_inputid.append(comment_tokens['input_ids']) |
| 615 | comments_mask.append(comment_tokens['attention_mask']) |
| 616 | comments_inputid = torch.LongTensor(comments_inputid) |
| 617 | comments_mask = torch.LongTensor(comments_mask) |
| 618 | |
| 619 | comments_like = [] |
| 620 | for num in item['comments_like']: |
| 621 | num_like = num.split(" ")[0] |
| 622 | comments_like.append(str2num(num_like)) |
| 623 | comments_like = torch.tensor(comments_like) |
| 624 | |
| 625 | return { |
| 626 | 'label': label, |
| 627 | 'comments_inputid': comments_inputid, |
| 628 | 'comments_mask': comments_mask, |
| 629 | 'comments_like': comments_like, |
| 630 | } |