active-learning-softmax-uncertainty-clipping/dataset_loader.py at btw_beta3 · jgonsior/active-learning-softmax-uncertainty-clipping · 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
80
81
from lib2to3.pgen2 import token
from transformers import AutoTokenizer
import datasets
import numpy as np
from small_text.integrations.transformers.datasets import TransformersDataset


def load_my_dataset(dataset: str, transformer_model_name: str, tokenization=True):
    if dataset == "ag_news":
        # works
        raw_dataset = datasets.load_dataset("ag_news")
    elif dataset == "trec6":
        # works
        raw_dataset = datasets.load_dataset("trec")
        raw_dataset = raw_dataset.rename_column("coarse_label", "label")
    elif dataset == "subj":
        # works
        raw_dataset = datasets.load_dataset("SetFit/subj")
    elif dataset == "rotten":
        # works
        raw_dataset = datasets.load_dataset("rotten_tomatoes")
    elif dataset == "imdb":
        # works
        raw_dataset = datasets.load_dataset("imdb")
    elif dataset == "sst2":
        raw_dataset = datasets.load_dataset("gpt3mix/sst2")
        #  raw_dataset = raw_dataset.rename_column("sentence", "text")

    elif dataset == "cola":
        raw_dataset = datasets.load_dataset("linxinyuan/cola")
        # raw_dataset = raw_dataset.rename_column("sentence", "text")
    else:
        print("dataset not known")
        exit(-1)

    print("First 3 training samples:\n")
    for i in range(3):
        print(raw_dataset["train"]["label"][i], " ", raw_dataset["train"]["text"][i])

    num_classes = np.unique(raw_dataset["train"]["label"]).shape[0]

    if tokenization:
        tokenizer = AutoTokenizer.from_pretrained(transformer_model_name)

        def _get_transformers_dataset(tokenizer, data, labels, max_length=60):
            data_out = []

            for i, doc in enumerate(data):
                encoded_dict = tokenizer.encode_plus(
                    doc,
                    add_special_tokens=True,
                    padding="max_length",
                    max_length=max_length,
                    return_attention_mask=True,
                    return_tensors="pt",
                    truncation="longest_first",
                )

                data_out.append(
                    (
                        encoded_dict["input_ids"],
                        encoded_dict["attention_mask"],
                        labels[i],
                    )
                )

            return TransformersDataset(data_out)

        # print(raw_dataset['train']['text'][:10])
        # print(raw_dataset['train']['label'][:10])

        train = _get_transformers_dataset(
            tokenizer, raw_dataset["train"]["text"], raw_dataset["train"]["label"]
        )
        test = _get_transformers_dataset(
            tokenizer, raw_dataset["test"]["text"], raw_dataset["test"]["label"]
        )

        return train, test, num_classes
    else:
        return raw_dataset["train"], raw_dataset["test"], num_classes