| import pandas as pd | |
| from transformers import AutoTokenizer | |
| class MyDataset: | |
| def __init__(self, data_file, tokenizer): | |
| self.data = pd.read_csv(data_file) | |
| self.tokenizer = tokenizer | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, idx): | |
| text = self.data.iloc[idx, 0] | |
| agents = self.data.iloc[idx, 1] | |
| actions = self.data.iloc[idx, 2] | |
| encoding = self.tokenizer.encode_plus( | |
| text, | |
| max_length=512, | |
| padding='max_length', | |
| truncation=True, | |
| return_attention_mask=True, | |
| return_tensors='pt' | |
| ) | |
| return { | |
| 'input_ids': encoding['input_ids'].flatten(), | |
| 'attention_mask': encoding['attention_mask'].flatten(), | |
| 'labels_agents': torch.tensor(agents), | |
| 'labels_actions': torch.tensor(actions) | |
| } |