methods/searchain/ColBERT/process_data.py

import json
import bz2
import os
import csv
import pandas as pd
import pickle
dataset = 'manyqa_text'

path = f'/data1/liuyaoyang/Papers/icml2025/datasets/{dataset}/corpus.pkl'

with open(path, "rb") as file:
    data = pickle.load(file)

def load_data(path=None):
    assert path
    corpus = []
    count = 0
    with open(
        f'/data1/liuyaoyang/Papers/icml2025/multi_rag/RAG/Search-in-the-Chain/data/{dataset}.tsv',
        'w', encoding='utf-8-sig') as w:
        writer = csv.writer(w, delimiter="\t")
        for item in list(data.values()):
            writer.writerow([count, item['text']])
            count += 1
            print(count)
    return corpus


load_data(path)