from collections import Counter
from tqdm import tqdm
import argparse
import csv
import pandas as pd
import os
import datasets

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--target', type=str, nargs='*')
    parser.add_argument('--output_dir', type=str, default=None)
    args = parser.parse_args()

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    for task in args.target:
        dataset = datasets.load_dataset('glue', task)
        if not os.path.exists(os.path.join(args.output_dir, task)):
            os.makedirs(os.path.join(args.output_dir, task))
        for split in list(dataset.keys()):
            dataset_split = dataset[split]
            entries = list(dataset_split[0].keys())
            with open(os.path.join(args.output_dir, task, "{}.csv".format(split)), 'w') as f:
                writer = csv.writer(f)
                writer.writerow(entries)
                for item in dataset_split:
                    
                    writer.writerow([item[ent] for ent in entries if ent != "idx"])

if __name__ == "__main__":
    main()