| 6 | import json |
| 7 | |
| 8 | class TextDataset(torch.utils.data.Dataset): |
| 9 | def __init__(self, json_file, text_key): |
| 10 | |
| 11 | self.data_list = json.load(open(json_file, 'r')) |
| 12 | self.data_list = [data for data in self.data_list if data[text_key] is not None] |
| 13 | self.text_key = text_key |
| 14 | |
| 15 | def __getitem__(self, data_id): |
| 16 | data = self.data_list[data_id] |
| 17 | text = data[self.text_key] |
| 18 | if isinstance(text, list): |
| 19 | text = text[0] |
| 20 | if 'sample_id' in data: |
| 21 | sample_id = data['sample_id'] |
| 22 | elif 'test_sample_id' in data: |
| 23 | sample_id = str(data['test_sample_id']) |
| 24 | elif 'global_id' in data: |
| 25 | sample_id = str(data['global_id']) |
| 26 | else: |
| 27 | raise ValueError("No sample_id or test_sample_id found in the data.") |
| 28 | |
| 29 | data_dict = {"text": text, "sample_id": sample_id} |
| 30 | |
| 31 | return data_dict |
| 32 | |
| 33 | def __len__(self): |
| 34 | return len(self.data_list) |
| 35 | |
| 36 | if __name__ == "__main__": |
| 37 |
no outgoing calls
no test coverage detected