docker

.gitignore

@@ -0,0 +1 @@
+*.tar.gz

api/.dockerignore

@@ -1,13 +1,11 @@
-bin
+__pycache__/
+bin/
+lib/
 build
-Dockerfile-alpine
+run
-lib
+*/__pycache__/
-__pycache__
+*/*/__pycache__/
-*/__pycache__
+*/*/*/__pycache__/
-*/*/__pycache__
-*/*/*/__pycache__
 pyvenv.cfg
 readme.md
-run
+chatbots_*
-test
-chatbots_api.tar.gz

api/Dockerfile

@@ -1,9 +1,7 @@
 # syntax-docker/dockerfile:1
-FROM python:slim
+FROM python:3.10-slim
-WORKDIR /app
+COPY . .
-COPY requirements.txt requirements.txt
 ENV PYTHONUNBUFFERED=1
 RUN pip install --no-cache-dir -r requirements.txt
-
+CMD [ "python3", "main.py"]
-COPY . .
+EXPOSE 6969
-CMD [ "uvicorn", "chatbots:api", "--proxy-headers", "--host", "0.0.0.0", "--port", "8000"]

api/build

@@ -1,5 +1,11 @@
 #!/bin/bash
 
-docker build --tag chatbots_api .
+FILE_PATH="../chatbots_api_$(uname -m).tar.gz"
-docker save chatbots_api | pigz > chatbots_api_$(uname -m).tar.gz
 
+if [ -f $FILE_PATH ]; then
+    rm $FILE_PATH
+fi
+
+#python -m nuitka --include-module=src --follow-imports --enable-plugin=torch main.py
+docker build --tag chatbots_api .
+docker save chatbots_api | pigz > $FILE_PATH

api/main.py

@@ -0,0 +1,6 @@
+import uvicorn
+import multiprocessing
+
+if __name__ == "__main__":
+    multiprocessing.freeze_support()
+    uvicorn.run('src.chatbots:api', host='0.0.0.0', port=6969, reload=True)

api/readme.md

@@ -1,9 +1,15 @@
 # Chatbots API
 
-FastAPI and PyTorch
+[FastAPI](https://fastapi.tiangolo.com/) and [PyTorch](https://pytorch.org/)
 
-To build yourself you'll need to first train a model with ../train
+To build one yourself you'll need to first [train a model](../train), 
+place the entire directory (checkpoints aren't needed) containing pytorch_model.bin in [bots](./src/bots), 
+then edit or duplicate [cartman.py](./src/bots/cartman.py).
 
-My image compressed is 2.1GB
+Cartman Docker images are availible for 
+[x86_64](https://doordesk.net/files/chatbots_api_x86_64.tar.gz) (1.6GB) and 
+[aarch64](https://doordesk.net/files/chatbots_api_x86_64.tar.gz) (1.4GB)
 
-Scripts in test to talk to it
+See [run](./run) and [test](./test) to interact with it
+
+Live demo [here](https://doordesk.net/cartman)

api/run

@@ -1,3 +1,7 @@
 #!/bin/bash
 
-docker run -d -p 6969:8000 chatbots_api
+source bin/activate &&
+python -m ensurepip &&
+pip install --upgrade -r requirements.txt &&
+clear &&
+python main.py

api/src/bots/cartman.py

@@ -4,10 +4,10 @@ from transformers.models.auto.tokenization_auto import AutoTokenizer
 from transformers.models.auto.modeling_auto import AutoModelForCausalLM
 
 tokenizer = AutoTokenizer.from_pretrained(
-    "microsoft/DialoGPT-large", padding_side='left'
+    "microsoft/DialoGPT-medium"
 )
 model = AutoModelForCausalLM.from_pretrained(
-    "src/bots/cartman"
+    "src/bots/cartman_med_3ep"
 )
 
 

api/src/chatbots.py

@@ -1,8 +1,8 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
-from src.models import Packet, BotResponse
+from .models import Packet, BotResponse
-from src.bots.cartman import cartman
+from .bots.cartman import cartman
 
 
 api = FastAPI()

train/.gitignore

@@ -1,3 +1,5 @@
 __pycache__/
 .ipynb_checkpoints/
 cartman/
+cached/
+runs/

train/clean.py

@@ -0,0 +1,28 @@
+import pandas as pd
+
+INPUT_FILE_PATH = './data/All-seasons.csv'
+OUPUT_FILE_PATH = './data/train_data.csv'
+
+
+df = pd.read_csv(INPUT_FILE_PATH)
+
+clean_lines = pd.Series(
+    [filter_lines
+        .replace('\n', '')
+        .replace('(', '')
+        .replace(')', '')
+        .replace('  ', ' ')
+        .strip()
+        for filter_lines in df.Line
+     ]
+)
+
+train_data = pd.DataFrame(df.Character)
+del df
+
+train_data['line'] = clean_lines
+
+train_data.columns = ['name', 'line']
+
+
+train_data.to_csv(OUPUT_FILE_PATH, index=False)

train/train.py

@@ -1,5 +1,20 @@
 # all the imports
 
+from transformers.models.auto.configuration_auto import AutoConfig
+from transformers.modeling_utils import PreTrainedModel
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+from transformers.models.auto.modeling_auto import (
+    AutoModelForCausalLM,
+    MODEL_WITH_LM_HEAD_MAPPING,
+)
+from transformers.utils import WEIGHTS_NAME
+from transformers.optimization import (
+    AdamW,
+    get_linear_schedule_with_warmup,
+)
+
+import torch
 import glob
 import logging
 import os
@@ -15,33 +30,22 @@ import pandas as pd
 from sklearn.model_selection import train_test_split
 
 from torch.nn.utils.rnn import pad_sequence
-from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler
+from torch.utils.data import (
-from torch.utils.data.distributed import DistributedSampler
+    DataLoader,
-from tqdm.notebook import tqdm, trange
+    Dataset,
-
+    RandomSampler,
-from pathlib import Path
+    SequentialSampler,
-
-from transformers import (
-    MODEL_WITH_LM_HEAD_MAPPING,
-    WEIGHTS_NAME,
-    AdamW,
-    AutoConfig,
-    PreTrainedModel,
-    PreTrainedTokenizer,
-    get_linear_schedule_with_warmup,
 )
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm, trange
 
-
+from torch.utils.tensorboard.writer import SummaryWriter
-try:
-    from torch.utils.tensorboard import SummaryWriter
-except ImportError:
-    from tensorboardX import SummaryWriter
 
 # --------------------------------------------------------------------------
 
 data = pd.read_csv('data/train.csv')
 
-CHARACTER_NAME = 'TARGET'
+CHARACTER_NAME = 'Cartman'
 contexted = []
 
 # context window of size 7
@@ -51,12 +55,12 @@ for i in data[data.name == CHARACTER_NAME].index:
     if i < n:
         continue
     row = []
-    prev = i - 1 - n # we additionally substract 1, so row will contain current response and 7 previous responses  
+    prev = i - 1 - n  # we additionally substract 1, so row will contain current response and 7 previous responses
     for j in range(i, prev, -1):
         row.append(data.line[j])
         contexted.append(row)
 
-columns = ['response', 'context'] 
+columns = ['response', 'context']
 columns = columns + ['context/' + str(i) for i in range(n - 1)]
 
 df = pd.DataFrame.from_records(contexted, columns=columns)
@@ -64,16 +68,21 @@ df = pd.DataFrame.from_records(contexted, columns=columns)
 trn_df, val_df = train_test_split(df, test_size=0.1)
 
 # create dataset suitable for our model
-def construct_conv(row, tokenizer, eos = True):
+
-    flatten = lambda l: [item for sublist in l for item in sublist]
+
-    conv = list(reversed([tokenizer.encode(x) + [tokenizer.eos_token_id] for x in row]))
+def construct_conv(row, tokenizer, eos=True):
+    def flatten(l): return [item for sublist in l for item in sublist]
+    conv = list(
+        reversed([tokenizer.encode(x) + [tokenizer.eos_token_id] for x in row]))
     conv = flatten(conv)
     return conv
 
+
 class ConversationDataset(Dataset):
     def __init__(self, tokenizer: PreTrainedTokenizer, args, df, block_size=512):
 
-        block_size = block_size - (tokenizer.model_max_length - tokenizer.max_len_single_sentence)
+        block_size = block_size - \
+            (tokenizer.model_max_length - tokenizer.max_len_single_sentence)
 
         directory = args.cache_dir
         cached_features_file = os.path.join(
@@ -81,7 +90,8 @@ class ConversationDataset(Dataset):
         )
 
         if os.path.exists(cached_features_file) and not args.overwrite_cache:
-            logger.info("Loading features from cached file %s", cached_features_file)
+            logger.info("Loading features from cached file %s",
+                        cached_features_file)
             with open(cached_features_file, "rb") as handle:
                 self.examples = pickle.load(handle)
         else:
@@ -92,9 +102,11 @@ class ConversationDataset(Dataset):
                 conv = construct_conv(row, tokenizer)
                 self.examples.append(conv)
 
-            logger.info("Saving features into cached file %s", cached_features_file)
+            logger.info("Saving features into cached file %s",
+                        cached_features_file)
             with open(cached_features_file, "wb") as handle:
-                pickle.dump(self.examples, handle, protocol=pickle.HIGHEST_PROTOCOL)
+                pickle.dump(self.examples, handle,
+                            protocol=pickle.HIGHEST_PROTOCOL)
 
     def __len__(self):
         return len(self.examples)
@@ -102,7 +114,8 @@ class ConversationDataset(Dataset):
     def __getitem__(self, item):
         return torch.tensor(self.examples[item], dtype=torch.long)
 
-# Cacheing and storing of data/checkpoints
+# Caching and storing of data/checkpoints
+
 
 def load_and_cache_examples(args, tokenizer, df_trn, df_val, evaluate=False):
     return ConversationDataset(tokenizer, args, df_val if evaluate else df_trn)
@@ -119,15 +132,18 @@ def set_seed(args):
 def _sorted_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) -> List[str]:
     ordering_and_checkpoint_path = []
 
-    glob_checkpoints = glob.glob(os.path.join(args.output_dir, "{}-*".format(checkpoint_prefix)))
+    glob_checkpoints = glob.glob(os.path.join(
+        args.output_dir, "{}-*".format(checkpoint_prefix)))
 
     for path in glob_checkpoints:
         if use_mtime:
             ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
         else:
-            regex_match = re.match(".*{}-([0-9]+)".format(checkpoint_prefix), path)
+            regex_match = re.match(
+                ".*{}-([0-9]+)".format(checkpoint_prefix), path)
             if regex_match and regex_match.groups():
-                ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
+                ordering_and_checkpoint_path.append(
+                    (int(regex_match.groups()[0]), path))
 
     checkpoints_sorted = sorted(ordering_and_checkpoint_path)
     checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
@@ -141,18 +157,19 @@ def _rotate_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) -
         return
 
     # Check if we should delete older checkpoint(s)
-    checkpoints_sorted = _sorted_checkpoints(args, checkpoint_prefix, use_mtime)
+    checkpoints_sorted = _sorted_checkpoints(
+        args, checkpoint_prefix, use_mtime)
     if len(checkpoints_sorted) <= args.save_total_limit:
         return
 
-    number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - args.save_total_limit)
+    number_of_checkpoints_to_delete = max(
+        0, len(checkpoints_sorted) - args.save_total_limit)
     checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
     for checkpoint in checkpoints_to_be_deleted:
-        logger.info("Deleting older checkpoint [{}] due to args.save_total_limit".format(checkpoint))
+        logger.info(
+            "Deleting older checkpoint [{}] due to args.save_total_limit".format(checkpoint))
         shutil.rmtree(checkpoint)
 
-from transformers import AutoModelWithLMHead, AutoModelForCausalLM, AutoTokenizer
-import torch
 
 tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
 model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
@@ -170,13 +187,15 @@ MODEL_CONFIG_CLASSES = list(MODEL_WITH_LM_HEAD_MAPPING.keys())
 MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
 
 # Args to allow for easy conversion of python script to notebook
+
+
 class Args():
     def __init__(self):
-        self.output_dir = 'models/output-medium'
+        self.output_dir = 'cartman/models/output-medium-3ep'
         self.model_type = 'gpt2'
-        self.model_name_or_path = 'microsoft/DialoGPT-large'
+        self.model_name_or_path = 'microsoft/DialoGPT-medium'
-        self.config_name = 'microsoft/DialoGPT-large'
+        self.config_name = 'microsoft/DialoGPT-medium'
-        self.tokenizer_name = 'microsoft/DialoGPT-large'
+        self.tokenizer_name = 'microsoft/DialoGPT-medium'
         self.cache_dir = 'cached'
         self.block_size = 512
         self.do_train = True
@@ -189,7 +208,7 @@ class Args():
         self.weight_decay = 0.0
         self.adam_epsilon = 1e-8
         self.max_grad_norm = 1.0
-        self.num_train_epochs = 4
+        self.num_train_epochs = 3
         self.max_steps = -1
         self.warmup_steps = 0
         self.logging_steps = 1000
@@ -205,8 +224,10 @@ class Args():
         self.fp16 = False
         self.fp16_opt_level = 'O1'
 
+
 args = Args()
 
+
 def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedTokenizer) -> Tuple[int, float]:
     """ Train the model """
     if args.local_rank in [-1, 0]:
@@ -219,22 +240,25 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
             return pad_sequence(examples, batch_first=True)
         return pad_sequence(examples, batch_first=True, padding_value=tokenizer.pad_token_id)
 
-    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
+    train_sampler = RandomSampler(
+        train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
     train_dataloader = DataLoader(
-        train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, collate_fn=collate, drop_last = True
+        train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, collate_fn=collate, drop_last=True
     )
 
     if args.max_steps > 0:
         t_total = args.max_steps
-        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
+        args.num_train_epochs = args.max_steps // (
+            len(train_dataloader) // args.gradient_accumulation_steps) + 1
     else:
-        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
+        t_total = len(
+            train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
 
-    model = model.module if hasattr(model, "module") else model  # Take care of distributed/parallel training
+    # Take care of distributed/parallel training
+    model = model.module if hasattr(model, "module") else model
     model.resize_token_embeddings(len(tokenizer))
     # add_special_tokens_(model, tokenizer)
 
-
     # Prepare optimizer and schedule (linear warmup and decay)
     no_decay = ["bias", "LayerNorm.weight"]
     optimizer_grouped_parameters = [
@@ -242,9 +266,11 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
             "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
             "weight_decay": args.weight_decay,
         },
-        {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0},
+        {"params": [p for n, p in model.named_parameters() if any(
+            nd in n for nd in no_decay)], "weight_decay": 0.0},
     ]
-    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
+    optimizer = AdamW(optimizer_grouped_parameters,
+                      lr=args.learning_rate, eps=args.adam_epsilon)
     scheduler = get_linear_schedule_with_warmup(
         optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
     )
@@ -256,15 +282,19 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
         and os.path.isfile(os.path.join(args.model_name_or_path, "scheduler.pt"))
     ):
         # Load in optimizer and scheduler states
-        optimizer.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
+        optimizer.load_state_dict(torch.load(
-        scheduler.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "scheduler.pt")))
+            os.path.join(args.model_name_or_path, "optimizer.pt")))
+        scheduler.load_state_dict(torch.load(
+            os.path.join(args.model_name_or_path, "scheduler.pt")))
 
     if args.fp16:
         try:
             from apex import amp
         except ImportError:
-            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
+            raise ImportError(
-        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)
+                "Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
+        model, optimizer = amp.initialize(
+            model, optimizer, opt_level=args.fp16_opt_level)
 
     # multi-gpu training (should be after apex fp16 initialization)
     if args.n_gpu > 1:
@@ -273,21 +303,24 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
     # Distributed training (should be after apex fp16 initialization)
     if args.local_rank != -1:
         model = torch.nn.parallel.DistributedDataParallel(
-            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True
+            model, device_ids=[
+                args.local_rank], output_device=args.local_rank, find_unused_parameters=True
         )
 
     # Train!
     logger.info("***** Running training *****")
     logger.info("  Num examples = %d", len(train_dataset))
     logger.info("  Num Epochs = %d", args.num_train_epochs)
-    logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
+    logger.info("  Instantaneous batch size per GPU = %d",
+                args.per_gpu_train_batch_size)
     logger.info(
         "  Total train batch size (w. parallel, distributed & accumulation) = %d",
         args.train_batch_size
         * args.gradient_accumulation_steps
         * (torch.distributed.get_world_size() if args.local_rank != -1 else 1),
     )
-    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
+    logger.info("  Gradient Accumulation steps = %d",
+                args.gradient_accumulation_steps)
     logger.info("  Total optimization steps = %d", t_total)
 
     global_step = 0
@@ -297,15 +330,21 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
     if args.model_name_or_path and os.path.exists(args.model_name_or_path):
         try:
             # set global_step to gobal_step of last saved checkpoint from model path
-            checkpoint_suffix = args.model_name_or_path.split("-")[-1].split("/")[0]
+            checkpoint_suffix = args.model_name_or_path.split(
+                "-")[-1].split("/")[0]
             global_step = int(checkpoint_suffix)
-            epochs_trained = global_step // (len(train_dataloader) // args.gradient_accumulation_steps)
+            epochs_trained = global_step // (len(train_dataloader) //
-            steps_trained_in_current_epoch = global_step % (len(train_dataloader) // args.gradient_accumulation_steps)
+                                             args.gradient_accumulation_steps)
+            steps_trained_in_current_epoch = global_step % (
+                len(train_dataloader) // args.gradient_accumulation_steps)
 
-            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
+            logger.info(
+                "  Continuing training from checkpoint, will skip to saved global_step")
             logger.info("  Continuing training from epoch %d", epochs_trained)
-            logger.info("  Continuing training from global step %d", global_step)
+            logger.info(
-            logger.info("  Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch)
+                "  Continuing training from global step %d", global_step)
+            logger.info("  Will skip the first %d steps in the first epoch",
+                        steps_trained_in_current_epoch)
         except ValueError:
             logger.info("  Starting fine-tuning.")
 
@@ -317,7 +356,8 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
     )
     set_seed(args)  # Added here for reproducibility
     for _ in train_iterator:
-        epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
+        epoch_iterator = tqdm(train_dataloader, desc="Iteration",
+                              disable=args.local_rank not in [-1, 0])
         for step, batch in enumerate(epoch_iterator):
 
             # Skip past any already trained steps if resuming training
@@ -326,12 +366,14 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
                 continue
 
             inputs, labels = (batch, batch)
-            if inputs.shape[1] > 1024: continue
+            if inputs.shape[1] > 1024:
+                continue
             inputs = inputs.to(args.device)
             labels = labels.to(args.device)
             model.train()
             outputs = model(inputs, labels=labels)
-            loss = outputs[0]  # model outputs are always tuple in transformers (see doc)
+            # model outputs are always tuple in transformers (see doc)
+            loss = outputs[0]
 
             if args.n_gpu > 1:
                 loss = loss.mean()  # mean() to average on multi-gpu parallel training
@@ -347,9 +389,11 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
             tr_loss += loss.item()
             if (step + 1) % args.gradient_accumulation_steps == 0:
                 if args.fp16:
-                    torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
+                    torch.nn.utils.clip_grad_norm_(
+                        amp.master_params(optimizer), args.max_grad_norm)
                 else:
-                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+                    torch.nn.utils.clip_grad_norm_(
+                        model.parameters(), args.max_grad_norm)
                 optimizer.step()
                 scheduler.step()  # Update learning rate schedule
                 model.zero_grad()
@@ -362,15 +406,19 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
                     ):  # Only evaluate when single GPU otherwise metrics may not average well
                         results = evaluate(args, model, tokenizer)
                         for key, value in results.items():
-                            tb_writer.add_scalar("eval_{}".format(key), value, global_step)
+                            tb_writer.add_scalar(
-                    tb_writer.add_scalar("lr", scheduler.get_lr()[0], global_step)
+                                "eval_{}".format(key), value, global_step)
-                    tb_writer.add_scalar("loss", (tr_loss - logging_loss) / args.logging_steps, global_step)
+                    tb_writer.add_scalar(
+                        "lr", scheduler.get_lr()[0], global_step)
+                    tb_writer.add_scalar(
+                        "loss", (tr_loss - logging_loss) / args.logging_steps, global_step)
                     logging_loss = tr_loss
 
                 if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
                     checkpoint_prefix = "checkpoint"
                     # Save model checkpoint
-                    output_dir = os.path.join(args.output_dir, "{}-{}".format(checkpoint_prefix, global_step))
+                    output_dir = os.path.join(
+                        args.output_dir, "{}-{}".format(checkpoint_prefix, global_step))
                     os.makedirs(output_dir, exist_ok=True)
                     model_to_save = (
                         model.module if hasattr(model, "module") else model
@@ -378,14 +426,18 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
                     model_to_save.save_pretrained(output_dir)
                     tokenizer.save_pretrained(output_dir)
 
-                    torch.save(args, os.path.join(output_dir, "training_args.bin"))
+                    torch.save(args, os.path.join(
+                        output_dir, "training_args.bin"))
                     logger.info("Saving model checkpoint to %s", output_dir)
 
                     _rotate_checkpoints(args, checkpoint_prefix)
 
-                    torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
+                    torch.save(optimizer.state_dict(), os.path.join(
-                    torch.save(scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
+                        output_dir, "optimizer.pt"))
-                    logger.info("Saving optimizer and scheduler states to %s", output_dir)
+                    torch.save(scheduler.state_dict(), os.path.join(
+                        output_dir, "scheduler.pt"))
+                    logger.info(
+                        "Saving optimizer and scheduler states to %s", output_dir)
 
             if args.max_steps > 0 and global_step > args.max_steps:
                 epoch_iterator.close()
@@ -401,11 +453,13 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
 
 # Evaluation of some model
 
+
 def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, df_trn, df_val, prefix="") -> Dict:
     # Loop to handle MNLI double evaluation (matched, mis-matched)
     eval_output_dir = args.output_dir
 
-    eval_dataset = load_and_cache_examples(args, tokenizer, df_trn, df_val, evaluate=True)
+    eval_dataset = load_and_cache_examples(
+        args, tokenizer, df_trn, df_val, evaluate=True)
     os.makedirs(eval_output_dir, exist_ok=True)
     args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
     # Note that DistributedSampler samples randomly
@@ -417,7 +471,7 @@ def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, df_tr
 
     eval_sampler = SequentialSampler(eval_dataset)
     eval_dataloader = DataLoader(
-        eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size, collate_fn=collate, drop_last = True
+        eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size, collate_fn=collate, drop_last=True
     )
 
     # multi-gpu evaluate
@@ -448,7 +502,8 @@ def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, df_tr
 
     result = {"perplexity": perplexity}
 
-    output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt")
+    output_eval_file = os.path.join(
+        eval_output_dir, prefix, "eval_results.txt")
     with open(output_eval_file, "w") as writer:
         logger.info("***** Eval results {} *****".format(prefix))
         for key in sorted(result.keys()):
@@ -459,13 +514,15 @@ def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, df_tr
 
 # Main runner
 
+
 def main(df_trn, df_val):
     args = Args()
-    
+
     if args.should_continue:
         sorted_checkpoints = _sorted_checkpoints(args)
         if len(sorted_checkpoints) == 0:
-            raise ValueError("Used --should_continue but no checkpoint was found in --output_dir.")
+            raise ValueError(
+                "Used --should_continue but no checkpoint was found in --output_dir.")
         else:
             args.model_name_or_path = sorted_checkpoints[-1]
 
@@ -505,8 +562,10 @@ def main(df_trn, df_val):
     # Set seed
     set_seed(args)
 
-    config = AutoConfig.from_pretrained(args.config_name, cache_dir=args.cache_dir)
+    config = AutoConfig.from_pretrained(
-    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, cache_dir=args.cache_dir)
+        args.config_name, cache_dir=args.cache_dir)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.tokenizer_name, cache_dir=args.cache_dir)
     model = AutoModelForCausalLM.from_pretrained(
         args.model_name_or_path,
         from_tf=False,
@@ -514,15 +573,17 @@ def main(df_trn, df_val):
         cache_dir=args.cache_dir,
     )
     model.to(args.device)
-    
+
     logger.info("Training/evaluation parameters %s", args)
 
     # Training
     if args.do_train:
-        train_dataset = load_and_cache_examples(args, tokenizer, df_trn, df_val, evaluate=False)
+        train_dataset = load_and_cache_examples(
+            args, tokenizer, df_trn, df_val, evaluate=False)
 
         global_step, tr_loss = train(args, train_dataset, model, tokenizer)
-        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)
+        logger.info(" global_step = %s, average loss = %s",
+                    global_step, tr_loss)
 
     # Saving best-practices: if you use save_pretrained for the model and tokenizer, you can reload them using from_pretrained()
     if args.do_train:
@@ -546,26 +607,4 @@ def main(df_trn, df_val):
         tokenizer = AutoTokenizer.from_pretrained(args.output_dir)
         model.to(args.device)
 
-    # Evaluation
-    results = {}
-    if args.do_eval and args.local_rank in [-1, 0]:
-        checkpoints = [args.output_dir]
-        if args.eval_all_checkpoints:
-            checkpoints = list(
-                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
-            )
-            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
-        logger.info("Evaluate the following checkpoints: %s", checkpoints)
-        for checkpoint in checkpoints:
-            global_step = checkpoint.split("-")[-1] if len(checkpoints) > 1 else ""
-            prefix = checkpoint.split("/")[-1] if checkpoint.find("checkpoint") != -1 else ""
-
-            model = AutoModelForCausalLM.from_pretrained(checkpoint)
-            model.to(args.device)
-            result = evaluate(args, model, tokenizer, df_trn, df_val, prefix=prefix)
-            result = dict((k + "_{}".format(global_step), v) for k, v in result.items())
-            results.update(result)
-
-    return results
-
 main(trn_df, val_df)