simpletransformers

1
import os
2
from datasets import load_dataset
3

4

5
os.makedirs("data/msmarco", exist_ok=True)
6

7
print("=== Downloading MSMARCO ===")
8
print("Downloading MSMARCO training triples...")
9
dataset = load_dataset("thilina/negative-sampling")["train"]
10

11
print("Dataset loaded. Sample:")
12
print(dataset[0])
13

14
qrels = load_dataset("BeIR/msmarco-qrels")["validation"]
15

16
print("Saving dataset to disk...")
17
# Save the dataset to disk
18
dataset.to_csv("data/msmarco/msmarco-train.tsv", sep="\t", index=False)
19
qrels.to_csv("data/msmarco/devs.tsv", sep="\t", index=False)
20

21
print("Done.")
22
print("=== MSMARCO download complete ===")
23