Merge CQADupStack subdatasets into a single dataset
(data_path: str, verbose: bool = False)
| 252 | |
| 253 | |
| 254 | def merge_cqa_dupstack(data_path: str, verbose: bool = False): |
| 255 | """Merge CQADupStack subdatasets into a single dataset""" |
| 256 | data_path = Path(data_path) |
| 257 | dataset = data_path.name |
| 258 | assert dataset == "cqadupstack", "Dataset must be CQADupStack" |
| 259 | |
| 260 | # check if corpus.jsonl exists |
| 261 | corpus_path = data_path / "corpus.jsonl" |
| 262 | if not corpus_path.exists(): |
| 263 | # combine all the corpus files into one |
| 264 | # corpus files are located under cqadupstack/<name>/corpus.jsonl |
| 265 | corpus_files = list(data_path.glob("*/corpus.jsonl")) |
| 266 | with open(corpus_path, "w") as f: |
| 267 | for file in tqdm(corpus_files, desc="Merging Corpus", disable=not verbose): |
| 268 | # get the name of the corpus |
| 269 | corpus_name = file.parent.name |
| 270 | |
| 271 | with open(file, "r") as f2: |
| 272 | for line in tqdm( |
| 273 | f2, |
| 274 | desc=f"Merging {corpus_name} Corpus", |
| 275 | leave=False, |
| 276 | disable=not verbose, |
| 277 | ): |
| 278 | # first, read with json |
| 279 | line = json.loads(line) |
| 280 | # add the corpus name to _id |
| 281 | line["_id"] = f"{corpus_name}_{line['_id']}" |
| 282 | # write back to file |
| 283 | f.write(json.dumps(line)) |
| 284 | f.write("\n") |
| 285 | |
| 286 | # now, do the same for queries.jsonl |
| 287 | queries_path = data_path / "queries.jsonl" |
| 288 | if not queries_path.exists(): |
| 289 | queries_files = list(data_path.glob("*/queries.jsonl")) |
| 290 | with open(queries_path, "w") as f: |
| 291 | for file in tqdm( |
| 292 | queries_files, desc="Merging Queries", disable=not verbose |
| 293 | ): |
| 294 | # get the name of the corpus |
| 295 | corpus_name = file.parent.name |
| 296 | |
| 297 | with open(file, "r") as f2: |
| 298 | for line in tqdm( |
| 299 | f2, |
| 300 | desc=f"Merging {corpus_name} Queries", |
| 301 | leave=False, |
| 302 | disable=not verbose, |
| 303 | ): |
| 304 | # first, read with json |
| 305 | line = json.loads(line) |
| 306 | # add the corpus name to _id |
| 307 | line["_id"] = f"{corpus_name}_{line['_id']}" |
| 308 | # write back to file |
| 309 | f.write(json.dumps(line)) |
| 310 | f.write("\n") |
| 311 |