add training and evaluation scripts

data/fine-tuning/create_online_reviews.py

@@ -0,0 +1,28 @@
+from datasets import load_dataset
+from transformers import AutoTokenizer
+import datasets
+import os
+
+
+
+def main():
+    datasets.set_caching_enabled(False)
+    tokenizer = AutoTokenizer.from_pretrained(r"/tokenizer/loc")
+    
+    tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+    
+    data_loc = "path/to/review/jsons"
+    data_files = [fil.path for fil in os.scandir(data_loc)]
+    dataset = load_dataset('online_reviews_loading.py', data_files=data_files)
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
+
+    def process_rating(example):
+        example["labels"] = [float(item) for item in example["rating"]]
+        return example
+    dataset = dataset["train"].map(tokenize_function, batched=True).map(process_rating,batched=True,remove_columns=['rating']).shuffle(seed=42).train_test_split(test_size=0.1)
+    
+
+if __name__ == "__main__":
+    main()

data/fine-tuning/create_xed.py

@@ -0,0 +1,57 @@
+
+
+import os
+import sys
+from datasets import load_dataset, concatenate_datasets
+from transformers import PreTrainedTokenizerFast
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+    default_data_collator,
+    GPT2Tokenizer
+)
+from transformers.trainer_utils import get_last_checkpoint
+from transformers import AutoModelWithLMHead, AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoModel
+
+from transformers import GPT2Model
+from transformers import GPT2TokenizerFast
+import transformers
+import torch
+import numpy as np
+import argparse
+
+tokenizer = AutoTokenizer.from_pretrained("/checkpoint/loc")
+tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+out_dir = "/out_dir/xed"
+max_length = 1024
+
+
+fi_annotated_raw = load_dataset("xed_en_fi","fi_annotated")
+fi_neutral_raw = load_dataset("xed_en_fi","fi_neutral")
+
+def to_arr(examples):
+    labels = []
+    for item in examples["labels"]:
+        labels.append([item])
+    return {"sentence":examples["sentence"],"labels":labels}
+fi_neutral_mapped = fi_neutral_raw["train"].map(to_arr, batched=True)
+
+fi_neutral_mapped_cast = fi_neutral_mapped.cast(fi_annotated_raw["train"].features)
+concat_raw_set = concatenate_datasets([fi_neutral_mapped_cast, fi_annotated_raw["train"]])#combine neutral and other labels into single dataset
+
+def tokenize_function(examples):
+    return tokenizer(examples["sentence"], padding="max_length", truncation=True, max_length=max_length)
+
+def to_arr_2(examples):
+    labels = []
+    for item in examples["labels"]:
+        label = np.zeros(9)
+        label[item] = 1
+        labels.append(label.tolist())
+    return {"sentence":examples["sentence"],"labels":labels}
+
+tokenized_datasets = concat_raw_set.map(tokenize_function, batched=True).map(to_arr_2, batched=True).shuffle(seed=42).train_test_split(test_size=0.1)
+tokenized_datasets.save_to_disk(out_dir)

data/fine-tuning/create_yle.py

@@ -0,0 +1,95 @@
+
+
+import os
+import json
+import sys
+from datasets import load_dataset, load_from_disk, concatenate_datasets, Dataset
+from transformers import PreTrainedTokenizerFast
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+    default_data_collator,
+)
+from transformers.trainer_utils import get_last_checkpoint
+from transformers import AutoModelWithLMHead, AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoModel
+
+from transformers import GPT2Model
+from transformers import GPT2TokenizerFast
+import transformers
+import torch
+import numpy as np
+
+root = r'G:\Data\yle\data'#download from kielipankki and extract
+
+texts = []
+subjects = []
+first_subjects = []
+first_ids = []
+subject_ids = []
+
+for path, subdirs, files in os.walk(root):
+    #Data is split into multiple files
+    for name in files:
+        print(os.path.join(path, name))
+        with open(os.path.join(path, name), encoding="utf8") as f:
+            data = json.load(f)
+
+            #Each file contains json with multiple articles
+            for i in range(len(data["data"])):
+                try:
+                    txt = ""
+                    s = [] #Subjects
+                    s_ids = []#Id for the subjects
+                    #From the content loop trough the content and get only heading as text as we do not want to add metadata to a text dataset
+                    for c in data["data"][i]["content"]:
+                        if c["type"] in ("heading","text"):
+                            txt += c["text"]
+                        txt += "\n"
+                    first = ""
+                    #An article contains n subjects. Loop trough those and also save which one was first. We want that as a distinct column in the dataset for performance.
+                    if "subjects" in data["data"][i]:#To know if we have a first subject, check first if we even have subjects in json.
+                        first = data["data"][i]["subjects"][0]["title"]["fi"]
+                        first_id = data["data"][i]["subjects"][0]["id"]
+                        for subject in data["data"][i]["subjects"]:
+                            s.append(subject["title"]["fi"])
+                            s_ids.append(subject["id"])
+                    first_subjects.append(first)
+                    first_ids.append(first_id)
+                    texts.append(txt)
+                    subjects.append(s)
+                    subject_ids.append(s_ids)
+                except:
+                    #Some texts contain formatting errors, just skip those as they are a neglible portion of all the articles.
+                    pass
+
+
+dataset = Dataset.from_dict({"text":texts, "subjects":subjects, "first_subject":first_subjects, "first_ids":first_ids, "subject_ids":subject_ids})
+
+tokenizer_loc = "/tokenizer_loc"
+
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_loc)
+tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+
+def find_major_subject(example):
+    good_subjects = ["urheilu","Kotimaan uutiset","Ulkomaat","jääkiekko","talous","politiikka","poliisi","Liikenne ja kuljetus","kulttuuri","puolueet","onnettomuudet","musiikki","Koulutus ja kasvatus","Venäjä","tieliikenne","luonto","autot","terveys","Helsinki","Pohjoismaat","kunnat","Eurooppa","rikokset","vaalit","Yhdysvallat","lainvalvonta"]
+    import numpy as np #Some scopes were broken on Windows so import again here to get batched processing to work...
+    example["main_subject"] = None
+    label = np.zeros(len(good_subjects))#sparse label matrix, to be made into one-hot later
+    for subject in example["subjects"]:
+        if subject in good_subjects:
+            example["main_subject"] = subject
+            label[good_subjects.index(subject)] = 1
+            #example["labels"] = label
+            break
+    return {"labels":label}
+
+filtered = dataset.map(find_major_subject, num_proc=12).filter(lambda example: example['main_subject'] != None)
+
+def tokenize_function(examples):
+    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=800)
+tokenized_and_filtered_dataset = filtered.map(tokenize_function, batched=True)
+
+tokenized_and_filtered_dataset.save_to_disk("/output/dir")

data/fine-tuning/create_ylilauta.py

@@ -0,0 +1,58 @@
+import os
+import sys
+from datasets import load_dataset, load_from_disk, concatenate_datasets, Dataset
+from transformers import PreTrainedTokenizerFast
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+    default_data_collator,
+)
+from transformers.trainer_utils import get_last_checkpoint
+from transformers import AutoModelWithLMHead, AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoModel
+
+from transformers import GPT2Model
+from transformers import GPT2TokenizerFast
+import transformers
+import torch
+import numpy as np
+import argparse
+
+tokenizer = AutoTokenizer.from_pretrained("/tokenizer/loc")
+tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+out_dir = "/out_dir/ylilauta"
+max_length = 1024
+
+#checkpoint_loc = r"H:\Data_temp\checkpoints\good_large\checkpoint-67400"
+#output_dir = r"H:\Data_temp\checkpoints\tests\yle"
+
+path = r"/data/ylilauta-corpus/data/100-percent/train.txt" #get from https://github.com/spyysalo/ylilauta-corpus
+text = []
+labels = []
+with open(path,"r",encoding="utf-8") as f:
+    for line in f:
+        parts = line.split(" ", maxsplit=1)
+        labels.append(parts[0])
+        text.append(parts[1])
+
+data_dict = {"text":text,"labels":labels}
+dataset = Dataset.from_dict(data_dict)
+label_names = dataset.unique('labels')
+n_labels = len(label_names)
+
+def to_one_hot(examples):
+    import numpy as np
+    
+    label = np.zeros(n_labels)
+    label[label_names.index(examples["labels"])] = 1
+
+    return {"text":examples["text"],"labels":label.tolist()}
+    
+def tokenize_function(examples):
+    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=max_length)
+
+tokenized = dataset.map(to_one_hot).map(tokenize_function).train_test_split(test_size=0.1)
+
+tokenized.save_to_disk(out_dir)

data/fine-tuning/online_reviews_loading.py

@@ -0,0 +1,51 @@
+from datasets import load_dataset
+import datasets
+import json
+import numpy as np
+import os
+#Dataset loading script that is missing quite a lot of details but works
+class NewDataset(datasets.GeneratorBasedBuilder):
+    def _info(self):
+        return datasets.DatasetInfo(
+            description="beep boop",
+            features=datasets.Features(
+                {
+                    "description": datasets.Value("string"),
+                    "text": datasets.Value("string"),
+                    "rating": datasets.Value("int32")
+                }
+            ),
+            # No default supervised_keys (as we have to pass both question
+            # and context as input).
+            supervised_keys=None,
+            homepage="no",
+            citation="no",
+        )
+        
+    def _split_generators(self, dl_manager):
+            files = self.config.data_files
+            return [datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"files": files["train"]})]
+    def _generate_examples(
+        self, files  # method parameters are unpacked from `gen_kwargs` as given in `_split_generators`
+    ):
+        """ Yields examples as (key, example) tuples. """
+        # This method handles input defined in _split_generators to yield (key, example) tuples from the dataset.
+        # The `key` is here for legacy reason (tfds) and is not important in itself.
+        #print("files",files)
+        key = 0
+        for file in files:
+            with open(file, encoding="utf-8") as f:
+                data = json.load(f)
+
+                for item in data:
+                    for review in item["reviews"]:
+                        yield key, {
+                            "description": item["description_raw"],
+                            "text": review["reviewText"],
+                            "rating": review["rating"],
+                        }
+                        key += 1
+
+
+
+                

data/tokenize.py

@@ -0,0 +1,50 @@
+from datasets import Dataset, load_dataset, concatenate_datasets
+import datasets
+from transformers import GPT2TokenizerFast
+from tokenizers.processors import TemplateProcessing
+
+input_dir = "dataset_location"
+tokenizer_file="path/to/file"
+output_dir="output/dir"
+tokenizer = GPT2TokenizerFast.from_pretrained(tokenizer_file)
+#Add eos tokens to the tokenization pipeline as they are not added otherwise
+tokenizer._tokenizer.post_processor = TemplateProcessing(
+    single="$0 "+tokenizer.eos_token,
+    pair="$A "+tokenizer.eos_token+" $B:1 "+tokenizer.eos_token,
+    special_tokens=[(tokenizer.eos_token, 0)],
+)
+
+def tokenize_function(examples):
+    return tokenizer(examples["text"])
+
+
+def group_texts(examples):
+    #group texts. This is based on Hugging Face CLM example
+    block_size = 1024
+    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+    total_len = len(concatenated_examples[list(examples.keys())[0]])
+    total_len = (total_len//block_size) * block_size
+    result = {
+        k: [t[i:i+block_size] for i in range(0, total_len, block_size)]
+        for k, t in concatenated_examples.items()
+    }
+    result["labels"] = result["input_ids"].copy()
+    return result
+
+def main():
+    num_proc=12 #set to something appropriate
+    dataset = datasets.load_from_disk(input_dir) #This one load a saved dataset object from disk. You could create a dataset from iterable or load one like:
+    #dataset = load_dataset("Finnish-NLP/mc4_fi_cleaned", split="train").remove_columns(["timestamp","url"]) #Example usage from Hugging Face Hub
+
+    #Tokenize, filter out very short texts and group texts to blocks of attention size
+    dataset\
+        .shuffle(seed=42, load_from_cache_file=False, writer_batch_size=100000)\
+        .map(tokenize_function, batched=True, num_proc=num_proc, remove_columns=dataset.column_names, load_from_cache_file=False, writer_batch_size=100000)\
+        .filter(lambda e: len(e["input_ids"]) > 20, num_proc=num_proc, load_from_cache_file=False, writer_batch_size=100000)\
+        .map(group_texts, batched=True, num_proc=num_proc, load_from_cache_file=False, writer_batch_size=100000)\
+        .train_test_split(test_size=0.05, load_from_cache_file=False, writer_batch_size=100000)\
+        .save_to_disk(output_dir)
+    print(dataset)
+
+if __name__ == "__main__":
+    main()

data/train_tokenizer.py

@@ -0,0 +1,25 @@
+from tokenizers import Tokenizer, normalizers, models, pre_tokenizers, processors, ByteLevelBPETokenizer
+import tokenizers
+
+from tokenizers.models import WordPiece, BPE
+from tokenizers.trainers import WordPieceTrainer, BpeTrainer
+from tokenizers.pre_tokenizers import Whitespace, Punctuation, Sequence
+from tokenizers.processors import TemplateProcessing
+import os
+
+from transformers import AutoTokenizer
+old_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+import datasets
+input_dir = "/dataset/location"
+dataset = datasets.load_from_disk(input_dir)
+
+def get_training_corpus():
+    for start_idx in range(0, len(dataset), 10000):
+        samples = dataset[start_idx : start_idx + 10000]
+        yield samples["text"]
+
+print("start")
+tokenizer = old_tokenizer.train_new_from_iterator(get_training_corpus(), vocab_size=50000)
+print("end")
+tokenizer.save_vocabulary("/tokenizer_location")

evaluate_and_analyze/evaluate.py

@@ -0,0 +1,132 @@
+
+
+import os
+import sys
+from datasets import load_dataset, load_from_disk, concatenate_datasets
+from transformers import PreTrainedTokenizerFast
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+    default_data_collator,
+)
+from transformers.trainer_utils import get_last_checkpoint
+from transformers import AutoModelWithLMHead, AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoModel
+
+from transformers import GPT2Model
+from transformers import GPT2TokenizerFast
+import transformers
+import torch
+import numpy as np
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('test', type=int)
+parser.add_argument('length', type=int)
+#parser.add_argument('--input_file', type=int)
+args = parser.parse_args()
+
+def compute_metrics(eval_pred):
+    logits,labels = eval_pred
+    import pickle
+    with open("logits_{}.pickle".format("xed"),"wb") as handle:
+        pickle.dump(logits, handle, protocol=pickle.HIGHEST_PROTOCOL)
+    with open("labels_{}.pickle".format("xed"),"wb") as handle:
+        pickle.dump(labels, handle, protocol=pickle.HIGHEST_PROTOCOL)
+    #Continue in a jupyter notebook from here
+
+    return
+
+
+
+class MultilabelTrainer(Trainer):
+    def compute_loss(self,model,inputs,return_outputs=False):
+        labels = inputs.pop("labels")
+        outputs = model(**inputs)
+        logits = outputs.logits
+        loss_fct = torch.nn.BCEWithLogitsLoss()
+        loss = loss_fct(logits.view(-1,self.model.config.num_labels),
+        labels.float().view(-1,self.model.config.num_labels))
+        return (loss,outputs) if return_outputs else loss
+
+def main():
+    ds_names = ["yle", "online_review","xed","ylilauta"]
+    #ds_sizes = [1000, 3000, 10000, 32000, 9999999]
+    print("test:",args.test)
+    ds_name = ds_names[args.test]
+    #ds_size = int(args.test.slit()[1])
+    ds_size = args.length
+    print(ds_name, ds_size)
+ 
+    metric = compute_metrics
+
+    #print("cuda_avail:",torch.cuda.is_available())
+    #checkpoint_loc = "/media/volume/output/checkpoint-275000"
+    #output_dir = "/media/volume/fi_nlp/output/finetune"
+    #checkpoint_loc = r"H:\Data_temp\checkpoints\good_large\checkpoint-67400"
+    output_dir = "/data/loc/"+ds_name
+
+    #Most of the parameters not used but lets just pass this to make the Trainer happy... 
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        learning_rate=5e-6,
+        adam_beta1=0.95,
+        adam_beta2=0.985,
+        adam_epsilon=1e-8,
+        weight_decay=0.001,
+        lr_scheduler_type="linear",
+        gradient_accumulation_steps=4,
+        max_steps=10000,
+        num_train_epochs=20000,
+        save_total_limit=2,
+        dataloader_num_workers=5,
+        save_steps=100000,
+        warmup_steps=500,
+        do_eval=True,
+        eval_steps=500,
+        evaluation_strategy="steps",
+        logging_strategy="steps",
+        logging_steps=50,
+        fp16_opt_level="O2",
+        half_precision_backend="amp",
+        log_on_each_node=False,
+        disable_tqdm=True
+    )
+
+    print(training_args)
+
+    dataset = load_from_disk(r"/data_loc/"+ds_name)["test"]
+    #dataset = load_from_disk(r"C:\Users\vin\Documents\Projects\dippa\tests\ylilauta\tokenized_set").train_test_split(test_size=0.1)
+    
+    trainer_class = MultilabelTrainer
+    
+    #print("num_labels",num_labels)
+    model = AutoModelForSequenceClassification.from_pretrained("/fine_tuning_checkpoint/"+ds_name)
+    tokenizer = AutoTokenizer.from_pretrained("/fine_tuning_checkpoint/"+ds_name)
+    tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+
+    print("init trainer")
+    trainer = trainer_class(
+            model=model,
+            args=training_args,
+            train_dataset=dataset,
+            eval_dataset=dataset,
+            tokenizer=tokenizer,
+            compute_metrics=metric,
+            data_collator=default_data_collator
+        )
+    #checkpoint = None
+    #checkpoint = get_last_checkpoint(output_dir)
+    #checkpoint = None
+    #train_result = trainer.train()
+    #trainer.save_state()
+    metrics = trainer.evaluate()
+    print(metrics)
+    #trainer.save_model()  # Saves the tokenizer too for easy upload
+
+if __name__ == "__main__":
+    main()

finetune.py

@@ -0,0 +1,150 @@
+
+
+import os
+import sys
+from datasets import load_dataset, load_from_disk, concatenate_datasets
+from transformers import PreTrainedTokenizerFast
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+    default_data_collator,
+)
+from transformers.trainer_utils import get_last_checkpoint
+from transformers import AutoModelWithLMHead, AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoModel
+
+from transformers import GPT2Model
+from transformers import GPT2TokenizerFast
+import transformers
+import torch
+import numpy as np
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('test', type=int)
+parser.add_argument('length', type=int)
+#parser.add_argument('--input_file', type=int)
+args = parser.parse_args()
+
+def compute_metrics(eval_pred):
+    logits,labels = eval_pred
+
+    predictions = np.zeros(logits.shape)
+    predictions[np.arange(len(predictions)),logits.argmax(1)] = 1
+    predictions = predictions > 0.5
+
+    #predictions = logits > 0.5
+    labels = labels > 0.5
+    return {"acc":np.all(predictions == labels,axis=1).sum()/predictions.shape[0]}
+
+def compute_metrics_regression(eval_pred):
+    logits,labels = eval_pred
+
+    labels = np.expand_dims(labels,1)
+    val = np.abs(logits-labels).mean()
+    perc = ((np.abs(logits-labels).round() < 1).sum()*100) / (len(labels))
+    perc_50 = ((np.abs(logits-labels).round()[0:50] < 1).sum()*100) / (50)
+
+    return {"dev":val,"perc":perc,"perc_50":perc_50}
+
+
+
+class MultilabelTrainer(Trainer):
+    def compute_loss(self,model,inputs,return_outputs=False):
+        labels = inputs.pop("labels")
+        outputs = model(**inputs)
+        logits = outputs.logits
+        loss_fct = torch.nn.BCEWithLogitsLoss()
+        loss = loss_fct(logits.view(-1,self.model.config.num_labels),
+        labels.float().view(-1,self.model.config.num_labels))
+        return (loss,outputs) if return_outputs else loss
+
+def main():
+    ds_names = ["yle", "online_reviews","xed","ylilauta"]
+    #ds_sizes = [1000, 3000, 10000, 32000, 9999999]
+    print("test:",args.test)
+    ds_name = ds_names[args.test]
+    ds_size = args.length
+    print(ds_name, ds_size)
+ 
+    metric = compute_metrics_regression if ds_name == "online_reviews" else compute_metrics
+
+    #print("cuda_avail:",torch.cuda.is_available())
+    #checkpoint_loc = "/media/volume/output/checkpoint-275000"
+    #output_dir = "/media/volume/fi_nlp/output/finetune"
+    #checkpoint_loc = r"H:\Data_temp\checkpoints\good_large\checkpoint-67400"
+    output_dir = "/scratch/project_462000007/hatanpav/output/dippa/gpt/"+ds_name
+
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        learning_rate=5e-6,
+        adam_beta1=0.95,
+        adam_beta2=0.985,
+        adam_epsilon=1e-8,
+        weight_decay=0.001,
+        lr_scheduler_type="linear",
+        gradient_accumulation_steps=2,#This one assumes 4x8 GPUs. Set to 64 to get global batch size of 64 with one GPU 
+        max_steps=10000,
+        num_train_epochs=20000,#Overriden by max_steps
+        save_total_limit=2,
+        dataloader_num_workers=5,
+        save_steps=100000,
+        warmup_steps=500,
+        do_eval=True,
+        eval_steps=500,
+        evaluation_strategy="steps",
+        logging_strategy="steps",
+        logging_steps=50,
+        fp16_opt_level="O2",
+        half_precision_backend="amp",
+        log_on_each_node=False,
+        disable_tqdm=True
+    )
+
+    print(training_args)
+
+    dataset = load_from_disk(r"/path/to/data/"+ds_name)
+    
+    #Handle regression type task:
+    n_labels = 1
+    trainer_class = MultilabelTrainer
+    try:
+        n_labels = len(dataset["train"][0]["labels"])
+    except:
+        #The case of label being a float.
+        n_labels = 1
+        trainer_class = Trainer
+    if ds_size > len(dataset["train"]):
+        ds_size = len(dataset["train"])
+    
+
+    model = AutoModelForSequenceClassification.from_pretrained("/checkpoint/loc",num_labels=n_labels)
+    tokenizer = AutoTokenizer.from_pretrained("/checkpoint/loc")
+    tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+
+    print("init trainer")
+    train_set = dataset["train"].select(range(ds_size))
+    test_set = dataset["test"]
+    trainer = trainer_class(
+            model=model,
+            args=training_args,
+            train_dataset=train_set,
+            eval_dataset=test_set,
+            tokenizer=tokenizer,
+            compute_metrics=metric,
+            data_collator=default_data_collator
+        )
+    checkpoint = None
+    #checkpoint = get_last_checkpoint(output_dir)
+    train_result = trainer.train(resume_from_checkpoint=checkpoint)
+    #trainer.save_state()
+    metrics = trainer.evaluate()
+    print(metrics)
+    trainer.save_model()  # Saves the tokenizer too for easy upload
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,97 @@
+
+import transformers
+import datasets
+from transformers import PreTrainedTokenizerFast
+from transformers import (
+    GPT2TokenizerFast,
+    AutoConfig,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+    default_data_collator
+)
+from transformers.trainer_utils import get_last_checkpoint
+import torch
+#from transformers.utils.dummy_tokenizers_objects import PreTrainedTokenizerFast
+
+#config_name = "C:\\Users\\vin\\Documents\\Projects\\NLP\\kielimalli\\config.json"
+#tokenizer_file = "C:\\Users\\vin\\Documents\\Projects\\NLP\\models\\tokens.json"
+#input_dir = "H:\\Data_temp\\tokenized_dataset"
+#output_dir = "H:\\Data_temp\\checkpoints\\model1"
+
+def main():
+    import os
+    #enable if required by your environment
+    #os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    #torch.backends.cuda.matmul.allow_tf32 = True
+    #torch.backends.cudnn.allow_tf32 = True
+    
+    config_name = "config_large_bpe.json"
+    tokenizer_files = "/path/to/tokenizer/files"
+    input_dir = "/data/dir"
+    output_dir = "/out/dir"
+
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        learning_rate=2.067e-5,
+        lr_scheduler_type="linear",
+        adam_beta1=0.95,
+        adam_beta2=0.985,
+        adam_epsilon=1e-8,
+        weight_decay=0.001,
+        gradient_accumulation_steps=32,
+        num_train_epochs=6.7,
+        save_total_limit=2,
+        dataloader_num_workers=10,
+        save_steps=100,
+        warmup_steps=1000,
+        do_eval=True,
+        eval_steps=1000,
+        evaluation_strategy="steps",
+        logging_strategy="steps",
+        logging_steps=100,
+        bf16=True,
+        tf32=True,
+        fp16_opt_level="O2",
+        half_precision_backend="amp",
+        bf16_full_eval=True
+    )
+
+    print("setting up tokenizer...")
+    tokenizer = GPT2TokenizerFast.from_pretrained(tokenizer_files)
+    #tokenizer.add_special_tokens({'pad_token': '[PAD]'})#Probably wrong
+    tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+    from tokenizers.processors import TemplateProcessing
+    tokenizer._tokenizer.post_processor = TemplateProcessing(
+        single="$0 "+tokenizer.eos_token,
+        pair="$A "+tokenizer.eos_token+" $B:1 "+tokenizer.eos_token,
+        special_tokens=[(tokenizer.eos_token, 0)],
+    )
+
+    print("loading model...")
+    config = AutoConfig.from_pretrained(config_name)
+    model = AutoModelForCausalLM.from_config(config)
+    #model = AutoModelForCausalLM.from_pretrained("/checkpoint/dir") if restarting training completely and loading weights from a checkpoints
+    model.gradient_checkpointing_enable() #Optional, affects performance
+    print("loading data...")
+    dataset = datasets.load_from_disk(input_dir)
+
+    print("starting training...")
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset["train"],
+        data_collator=default_data_collator,
+        eval_dataset=dataset["test"].select(range(10000)), #To save time do not evaluate on whole test set during training
+        tokenizer=tokenizer
+    )
+
+    #checkpoint = None
+    checkpoint = get_last_checkpoint(output_dir)
+    print("checkpoint:", checkpoint)
+    trainer.train(resume_from_checkpoint=checkpoint)
+
+if __name__ == "__main__":
+    main()