| 145 | |
| 146 | |
| 147 | class SidTokenFeatDataset(Dataset): |
| 148 | def __init__(self, description_file, tokenizer=None, max_len=2048, sample=-1, test=False, seed=0, category=""): |
| 149 | """ |
| 150 | Dataset for token2description and description2token tasks. |
| 151 | |
| 152 | Args: |
| 153 | description_file: Path to .index.json file with token feature |
| 154 | tokenizer: Tokenizer for encoding text |
| 155 | max_len: Maximum sequence length |
| 156 | sample: Number of samples to use (-1 for all) |
| 157 | test: Whether this is test mode |
| 158 | seed: Random seed |
| 159 | category: Category name for prompts |
| 160 | """ |
| 161 | random.seed(seed) |
| 162 | |
| 163 | # Load item features and indices |
| 164 | with open(description_file, 'r') as f: |
| 165 | self.token_feat = json.load(f) |
| 166 | |
| 167 | self.tokenizer = Tokenizer(tokenizer) if tokenizer is not None else None |
| 168 | self.test = test |
| 169 | self.max_len = max_len |
| 170 | self.category = category |
| 171 | |
| 172 | # Build sid2title and title2sid mappings |
| 173 | self.token2description = {} |
| 174 | self.description2token = {} |
| 175 | |
| 176 | for example in self.token_feat: |
| 177 | prefix = example['token'] |
| 178 | description = example['description'] |
| 179 | self.token2description[prefix] = description |
| 180 | self.description2token[description] = prefix |
| 181 | |
| 182 | # Create data samples |
| 183 | self.data = [] |
| 184 | |
| 185 | # Create sid2title samples |
| 186 | for prefix, description in self.token2description.items(): |
| 187 | self.data.append({ |
| 188 | 'task': 'token2description', |
| 189 | 'input': prefix, |
| 190 | 'output': description |
| 191 | }) |
| 192 | |
| 193 | # Create title2sid samples |
| 194 | for description, prefix in self.description2token.items(): |
| 195 | self.data.append({ |
| 196 | 'task': 'description2token', |
| 197 | 'input': description, |
| 198 | 'output': prefix |
| 199 | }) |
| 200 | |
| 201 | if sample > 0 and sample < len(self.data): |
| 202 | self.data = random.sample(self.data, sample) |
| 203 | |
| 204 | if self.tokenizer is not None: |