-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdata_loader_prompt.py
64 lines (46 loc) · 2.89 KB
/
data_loader_prompt.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
from torch.utils.data import Dataset
import random
class Dataload_prompt(Dataset):
def __init__(self, df, df_dem, task):
self.df = df
self.df_dem = df_dem
self.task = task
self.turn = df['Turn'].values
self.label = df['type_ordinal'].values
self.prompts = df['Type'].values
self.index = df.index.values
def __getitem__(self, item):
turn = self.turn[item]
label = self.label[item]
prompts = self.prompts[item]
if self.task == 'train':
fluent = self.df[(self.df['type_ordinal']==0) & (self.df['Turn']!=turn)].sample(1)
fluent_turn = fluent['Turn'].to_list()[0]
fluent_prompt = fluent['Type'].to_list()[0]
anomia = self.df[(self.df['type_ordinal']==1) & (self.df['Turn']!=turn)].sample(1)
anomia_turn = anomia['Turn'].to_list()[0]
anomia_prompt = anomia['Type'].to_list()[0]
disf = self.df[(self.df['type_ordinal']==2) & (self.df['Turn']!=turn)].sample(1)
disf_turn = disf['Turn'].to_list()[0]
disf_prompt = disf['Type'].to_list()[0]
gram = self.df[(self.df['type_ordinal']==3) & (self.df['Turn']!=turn)].sample(1)
gram_turn = gram['Turn'].to_list()[0]
gram_prompt = gram['Type'].to_list()[0]
turn = turn + '. It is <mask>. ' + fluent_turn + '. It is ' + fluent_prompt + '. ' + anomia_turn + '. It is ' + anomia_prompt + '. ' + disf_turn + '. It is ' + disf_prompt + '. ' + gram_turn + '. It is ' + gram_prompt + '.'
else:
fluent = self.df_dem[(self.df_dem['type_ordinal']==0) & (self.df_dem['Turn']!=turn)].sample(1)
fluent_turn = fluent['Turn'].to_list()[0]
fluent_prompt = fluent['Type'].to_list()[0]
anomia = self.df_dem[(self.df_dem['type_ordinal']==1) & (self.df_dem['Turn']!=turn)].sample(1)
anomia_turn = anomia['Turn'].to_list()[0]
anomia_prompt = anomia['Type'].to_list()[0]
disf = self.df_dem[(self.df_dem['type_ordinal']==2) & (self.df_dem['Turn']!=turn)].sample(1)
disf_turn = disf['Turn'].to_list()[0]
disf_prompt = disf['Type'].to_list()[0]
gram = self.df_dem[(self.df_dem['type_ordinal']==3) & (self.df_dem['Turn']!=turn)].sample(1)
gram_turn = gram['Turn'].to_list()[0]
gram_prompt = gram['Type'].to_list()[0]
turn = turn + '. It is <mask>. ' + fluent_turn + '. It is ' + fluent_prompt + '. ' + anomia_turn + '. It is ' + anomia_prompt + '. ' + disf_turn + '. It is ' + disf_prompt + '. ' + gram_turn + '. It is ' + gram_prompt + '.'
return (turn, label, prompts)
def __len__ (self):
return len(self.turn)# -*- coding: utf-8 -*-