LexText/prepare_data.py at main · MeLeLBGU/LexText · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

import os
import json
from nltk.tokenize import word_tokenize
import tqdm
import random
from datasets import load_dataset
import pandas as pd

def load_data(dataset):
    # convert dataset
        # this isn't really needed...
        print("Loading dataset", dataset)
        if not isinstance(dataset, str):
            train_df = dataset.to_pandas() #doc["ctx_a"] + " " + doc["ctx_b"].capitalize()
            return train_df, None
        if "sst2" in dataset:
            data = load_dataset(dataset)
            text_key = "sentence"
            dataset_splits = ["train", "validation"]
        elif "imdb" in dataset:
            data = load_dataset(dataset)
            text_key = "text"
            dataset_splits = ["train", "test"]
        elif "qnli" in dataset:
            type = "validation"
            ds = load_dataset("glue", dataset)
            texts = []
            labels = []
            dataset_splits = ["train", "validation"]
        elif "mrpc" in dataset:
            ds = load_dataset("glue", "mrpc")
            dataset_splits = ["train", "validation"]
        elif "cola" in dataset:
            data = load_dataset("glue", "cola")
            text_key = "sentence"
            dataset_splits = ["train", "validation"]
        elif "ag_news" in dataset:
            ds = load_dataset("fancyzhx/ag_news")
            text_key = "text"
            dataset_splits = ["train", "test"]

        for dataset_split in dataset_splits:
            if "qnli" in dataset or "mrpc" in dataset:
                texts = []
                labels = []
                data = ds[dataset_split]
                if "qnli" in dataset:
                    for q, s, l in zip(data["question"], data["sentence"], data["label"]):
                        texts.append(q + "\t" + s)
                        labels.append(l)
                else:  # mrpc
                    for s1, s2, l in zip(data["sentence1"], data["sentence2"], data["label"]):
                        texts.append(s1 + "\t" + s2)
                        labels.append(l)
            else:
                texts = [sent[text_key] for sent in data[dataset_split]]
                labels = [sent["label"] for sent in data[dataset_split]]

            ids = [i for i in range(len(labels))]

            data_df = pd.DataFrame({
                'id': ids,
                'sentence': texts,
                'label': labels
            })
            if "sst2" and dataset_split == "validation":
                dataset_split = "test"
            output_file = f"./lexicalcustext/datasets/{dataset}/{dataset_split}.tsv"
            data_df.to_csv(output_file, index=False, sep='\t')
        #print(f"./datasets/{dataset}/train.tsv")
tasks = ["qnli", "sst2", "mrpc", "cola"]

for task in tasks:
    os.makedirs(f"lexicalcustext/theirs_{task}", exist_ok=True)
    os.makedirs(f"custext/theirs_{task}", exist_ok=True)
    load_data(task)