| 157 | from torch.utils.data import Dataset |
| 158 | |
| 159 | class NamesDataset(Dataset): |
| 160 | |
| 161 | def __init__(self, data_dir): |
| 162 | self.data_dir = data_dir #for provenance of the dataset |
| 163 | self.load_time = time.localtime #for provenance of the dataset |
| 164 | labels_set = set() #set of all classes |
| 165 | |
| 166 | self.data = [] |
| 167 | self.data_tensors = [] |
| 168 | self.labels = [] |
| 169 | self.labels_tensors = [] |
| 170 | |
| 171 | #read all the ``.txt`` files in the specified directory |
| 172 | text_files = glob.glob(os.path.join(data_dir, '*.txt')) |
| 173 | for filename in text_files: |
| 174 | label = os.path.splitext(os.path.basename(filename))[0] |
| 175 | labels_set.add(label) |
| 176 | lines = open(filename, encoding='utf-8').read().strip().split('\n') |
| 177 | for name in lines: |
| 178 | self.data.append(name) |
| 179 | self.data_tensors.append(lineToTensor(name)) |
| 180 | self.labels.append(label) |
| 181 | |
| 182 | #Cache the tensor representation of the labels |
| 183 | self.labels_uniq = list(labels_set) |
| 184 | for idx in range(len(self.labels)): |
| 185 | temp_tensor = torch.tensor([self.labels_uniq.index(self.labels[idx])], dtype=torch.long) |
| 186 | self.labels_tensors.append(temp_tensor) |
| 187 | |
| 188 | def __len__(self): |
| 189 | return len(self.data) |
| 190 | |
| 191 | def __getitem__(self, idx): |
| 192 | data_item = self.data[idx] |
| 193 | data_label = self.labels[idx] |
| 194 | data_tensor = self.data_tensors[idx] |
| 195 | label_tensor = self.labels_tensors[idx] |
| 196 | |
| 197 | return label_tensor, data_tensor, data_label, data_item |
| 198 | |
| 199 | |
| 200 | ######################### |
no outgoing calls
no test coverage detected