.gitignore

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

api/.dockerignore

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

api/Dockerfile

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

api/build

@@ -1,11 +1,5 @@
 #!/bin/bash
 
-FILE_PATH="../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
+docker save chatbots_api | pigz > chatbots_api_$(uname -m).tar.gz
+

api/chatbots.py

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

api/main.py

@@ -1,6 +0,0 @@
-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,15 +1,9 @@
 # Chatbots API
 
-[FastAPI](https://fastapi.tiangolo.com/) and [PyTorch](https://pytorch.org/)
+FastAPI and PyTorch
 
-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).
+To build yourself you'll need to first train a model with ../train
 
-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)
+My image compressed is 2.1GB
 
-See [run](./run) and [test](./test) to interact with it
-
-Live demo [here](https://doordesk.net/cartman)
+Scripts in test to talk to it

api/run

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

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-medium"
+    "microsoft/DialoGPT-large", padding_side='left'
 )
 model = AutoModelForCausalLM.from_pretrained(
-    "src/bots/cartman_med_3ep"
+    "src/bots/cartman"
 )
 
 

train/.gitignore

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

train/clean.py

@@ -1,28 +0,0 @@
-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,20 +1,5 @@
 # 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
@@ -30,22 +15,33 @@ 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 DataLoader, Dataset, RandomSampler, SequentialSampler
 from torch.utils.data.distributed import DistributedSampler
-from tqdm import tqdm, trange
+from tqdm.notebook import tqdm, trange
 
-from torch.utils.tensorboard.writer import SummaryWriter
+from pathlib import Path
+
+from transformers import (
+    MODEL_WITH_LM_HEAD_MAPPING,
+    WEIGHTS_NAME,
+    AdamW,
+    AutoConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    get_linear_schedule_with_warmup,
+)
+
+
+try:
+    from torch.utils.tensorboard import SummaryWriter
+except ImportError:
+    from tensorboardX import SummaryWriter
 
 # --------------------------------------------------------------------------
 
 data = pd.read_csv('data/train.csv')
 
-CHARACTER_NAME = 'Cartman'
+CHARACTER_NAME = 'TARGET'
 contexted = []
 
 # context window of size 7
@@ -55,12 +51,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)
@@ -68,21 +64,16 @@ 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):
-    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]))
+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]))
     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(
@@ -90,8 +81,7 @@ 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:
@@ -102,11 +92,9 @@ 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)
@@ -114,8 +102,7 @@ class ConversationDataset(Dataset):
     def __getitem__(self, item):
         return torch.tensor(self.examples[item], dtype=torch.long)
 
-# Caching and storing of data/checkpoints
-
+# Cacheing 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)
@@ -132,18 +119,15 @@ 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]
@@ -157,19 +141,18 @@ 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")
@@ -187,15 +170,13 @@ 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 = 'cartman/models/output-medium-3ep'
+        self.output_dir = 'models/output-medium'
         self.model_type = 'gpt2'
-        self.model_name_or_path = 'microsoft/DialoGPT-medium'
-        self.config_name = 'microsoft/DialoGPT-medium'
-        self.tokenizer_name = 'microsoft/DialoGPT-medium'
+        self.model_name_or_path = 'microsoft/DialoGPT-large'
+        self.config_name = 'microsoft/DialoGPT-large'
+        self.tokenizer_name = 'microsoft/DialoGPT-large'
         self.cache_dir = 'cached'
         self.block_size = 512
         self.do_train = True
@@ -208,7 +189,7 @@ class Args():
         self.weight_decay = 0.0
         self.adam_epsilon = 1e-8
         self.max_grad_norm = 1.0
-        self.num_train_epochs = 3
+        self.num_train_epochs = 4
         self.max_steps = -1
         self.warmup_steps = 0
         self.logging_steps = 1000
@@ -224,10 +205,8 @@ 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]:
@@ -240,25 +219,22 @@ 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
 
-    # Take care of distributed/parallel training
-    model = model.module if hasattr(model, "module") else model
+    model = model.module if hasattr(model, "module") else model  # Take care of distributed/parallel training
     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 = [
@@ -266,11 +242,9 @@ 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
     )
@@ -282,19 +256,15 @@ 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")))
-        scheduler.load_state_dict(torch.load(
-            os.path.join(args.model_name_or_path, "scheduler.pt")))
+        optimizer.load_state_dict(torch.load(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.")
-        model, optimizer = amp.initialize(
-            model, optimizer, opt_level=args.fp16_opt_level)
+            raise ImportError("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:
@@ -303,24 +273,21 @@ 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
@@ -330,21 +297,15 @@ 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)
-            steps_trained_in_current_epoch = global_step % (
-                len(train_dataloader) // args.gradient_accumulation_steps)
+            epochs_trained = global_step // (len(train_dataloader) // 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("  Will skip the first %d steps in the first epoch",
-                        steps_trained_in_current_epoch)
+            logger.info("  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.")
 
@@ -356,8 +317,7 @@ 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
@@ -366,14 +326,12 @@ 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)
-            # model outputs are always tuple in transformers (see doc)
-            loss = outputs[0]
+            loss = outputs[0]  # model outputs are always tuple in transformers (see doc)
 
             if args.n_gpu > 1:
                 loss = loss.mean()  # mean() to average on multi-gpu parallel training
@@ -389,11 +347,9 @@ 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()
@@ -406,19 +362,15 @@ 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(
-                        "lr", scheduler.get_lr()[0], global_step)
-                    tb_writer.add_scalar(
-                        "loss", (tr_loss - logging_loss) / args.logging_steps, global_step)
+                            tb_writer.add_scalar("eval_{}".format(key), value, 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
@@ -426,18 +378,14 @@ 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(scheduler.state_dict(), os.path.join(
-                        output_dir, "scheduler.pt"))
-                    logger.info(
-                        "Saving optimizer and scheduler states to %s", output_dir)
+                    torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
+                    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()
@@ -453,13 +401,11 @@ 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
@@ -471,7 +417,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
@@ -502,8 +448,7 @@ 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()):
@@ -514,15 +459,13 @@ 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]
 
@@ -562,10 +505,8 @@ def main(df_trn, df_val):
     # Set seed
     set_seed(args)
 
-    config = AutoConfig.from_pretrained(
-        args.config_name, cache_dir=args.cache_dir)
-    tokenizer = AutoTokenizer.from_pretrained(
-        args.tokenizer_name, cache_dir=args.cache_dir)
+    config = AutoConfig.from_pretrained(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,
@@ -573,17 +514,15 @@ 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:
@@ -607,4 +546,26 @@ 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)