add personalized-chat-bot

personalized-chat-bot/util/bloom_trainer.py

@@ -0,0 +1,91 @@
+import sys
+
+import numpy as np
+from torch.utils.data import DataLoader
+from torch.optim import AdamW
+from transformers import get_scheduler
+import torch
+
+
+from util.metrics import perplexity
+
+
+class BloomTrainer:
+    DEFAULT_VAL_FREQ = 5
+    ITERATION_LIMIT = 150
+
+    def __init__(self, model, config, train_dataset, val_dataset, wandb_run=None, prompt_path=None, val_freq=None):
+        self.model = model
+        self.config = config
+        self.train_dataset = train_dataset
+        self.val_dataset = val_dataset
+        self.wandb_run = wandb_run
+        self.val_freq = val_freq
+        if self.val_freq is None:
+            self.val_freq = self.DEFAULT_VAL_FREQ
+        self.prompt_path = prompt_path
+
+        self.best_loss = np.inf
+
+        self.train_loader = DataLoader(self.train_dataset,
+                                       shuffle=True,
+                                       batch_size=config.BATCH_SIZE,
+                                       drop_last=True)
+        self.val_loader = DataLoader(self.val_dataset,
+                                     shuffle=True,
+                                     batch_size=config.BATCH_SIZE,
+                                     drop_last=False)
+
+        self.optimizer = AdamW(self.model.parameters(), lr=config.LR, weight_decay=config.WEIGHT_DECAY)
+
+        self.lr_scheduler = get_scheduler(
+            name="linear",
+            optimizer=self.optimizer,
+            num_warmup_steps=0,
+            num_training_steps= len(self.train_loader) * self.config.N_EPOCH
+        )
+
+    def train(self):
+        self.model.train()
+        iter_counter = 0
+        for epoch in range(self.config.N_EPOCH):
+            for batch in self.train_loader:
+                batch = {'input_ids': torch.stack(batch['input_ids']).T.to(self.config.DEVICE),
+                         'labels': torch.stack(batch['labels']).T.to(self.config.DEVICE)}
+                outputs = self.model(**batch)
+                loss = outputs.loss
+                loss.backward()
+                self.optimizer.step()
+                self.lr_scheduler.step()
+                self.optimizer.zero_grad()
+                self.wandb_run.log({'loss': loss})
+                iter_counter += 1
+                if (iter_counter + 1) % self.val_freq == 0:
+                    eval_perplexity = self.evaluate(perplexity)
+                    self.wandb_run.log({'perplexity': eval_perplexity})
+                    if loss.item() < self.best_loss:
+                        self.best_loss = loss.item()
+                        self.save_model(self.prompt_path)
+                        print('Model saved')
+                if iter_counter >= self.ITERATION_LIMIT:
+                    return
+
+    def evaluate(self, eval_fn):
+        logits = []
+        labels = []
+        self.model.eval()
+        with torch.no_grad():
+            for batch in self.val_loader:
+                batch = {'input_ids': torch.stack(batch['input_ids']).T.to(self.config.DEVICE),
+                         'labels': torch.stack(batch['labels']).T.to(self.config.DEVICE)}
+                outputs = self.model(**batch)
+                labels.extend(batch['input_ids'])
+                logits.extend(outputs.logits)
+        metric = eval_fn(logits, labels)
+        return metric
+
+    def save_model(self, path):
+        torch.save(self.model.transformer.prompt_embeddings.state_dict(), path)
+
+    def load_model(self, path):
+        self.model.transformer.prompt_embeddings.load_state_dict(torch.load(path))

personalized-chat-bot/util/data.py

@@ -0,0 +1,74 @@
+import itertools
+from torch.utils.data import Dataset
+import numpy as np
+from joblib import Parallel, delayed
+
+
+class OnePersonaDataset(Dataset):
+    def __init__(self, data, tokenizer, transforms=None, positive_candidates=True, n_jobs=8):
+        super().__init__()
+
+        self.data = data
+        if len(data) == 0:
+            self.input_ids = []
+            self.history = []
+            self.labels = []
+            return
+
+        if positive_candidates:
+            self.history = [row['history'] + [row['candidates'][-1], ] for row in data]
+            self.labels = np.ones(len(self.history), dtype=int)
+        else:
+            self.history = [row['history'] + [candidate, ] for row in data
+                            for candidate in row['candidates']]
+            self.labels = itertools.chain.from_iterable([0] * (len(row['candidates']) - 1) + [1]
+                                                        for row in data)
+            self.labels = np.array(self.labels, dtype=int)
+
+        if transforms is None:
+            self.history = ["\n".join(item) for item in self.history]
+        else:
+            self.history = Parallel(n_jobs=n_jobs)(delayed(transforms)(item) for item in self.history)
+        self.input_ids = tokenizer(self.history, padding='max_length', truncation=True)["input_ids"]
+
+    def __getitem__(self, idx):
+        return {'input_ids': self.input_ids[idx],
+                'labels': self.input_ids[idx],
+                'example': self.history[idx],
+                'class': self.labels[idx]}
+
+    def __len__(self):
+        return len(self.data)
+
+
+class PersonaChatDataset(Dataset):
+    DEFAULT_DATASET_NAME = "bavard/personachat_truecased"
+
+    def __init__(self, clustering, dataset, tokenizer):
+        super().__init__()
+
+        self.dataset = dataset
+        self.clustering = clustering
+
+        all_personalities = list(set([sent for item in self.dataset
+                                      for sent in item['personality']]))
+        predicted_centers = self.clustering.predict(all_personalities)
+        self.all_personalities_to_id = {persona: center
+                                        for persona, center in zip(all_personalities, predicted_centers)}
+        self.personalities = self.clustering._cluster_centers
+
+        subdataset_data_by_personality = [[] for _ in range(len(self.personalities))]
+
+        for i in range(len(self.dataset)):
+            item = self.dataset[i]
+            cur_persona_ids = [self.all_personalities_to_id[persona] for persona in item['personality']]
+            for persona_id in cur_persona_ids:
+                subdataset_data_by_personality[persona_id].append(item)
+
+        self.subdatasets = [OnePersonaDataset(cur_data, tokenizer) for cur_data in subdataset_data_by_personality]
+
+    def __getitem__(self, persona_id):
+        return self.subdatasets[persona_id]
+
+    def __len__(self, ):
+        return len(self.datasets)

personalized-chat-bot/util/dialogue_manager.py

@@ -0,0 +1,27 @@
+from transformers import DistilBertForSequenceClassification
+from torch import nn
+
+class DialogueManagerModel(nn.Module):
+    DEFAULT_MODEL = "distilbert-base-uncased"
+
+    def __init__(self, n_classes, model_name=None, device='cpu'):
+        super().__init__()
+        if model_name is None:
+            self.model = DistilBertForSequenceClassification.from_pretrained(self.DEFAULT_MODEL)
+        else:
+            raise NotImplementedError()
+        self.model.to(device)
+        self.n_classes = n_classes
+        self.freeze_layers()
+        self.model.classifier = nn.Linear(self.model.classifier.in_features, self.n_classes,
+                                          device=device)
+
+        for param in self.model.classifier.parameters():
+            param.requires_grad = True
+
+    def freeze_layers(self):
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def forward(self, X):
+        return self.model(X)

personalized-chat-bot/util/metrics.py

@@ -0,0 +1,27 @@
+import numpy as np
+import scipy
+import torch
+
+
+def _perplexity(logits, labels, pad_token=3):
+    for i in range(len(labels)-1, -1, -1):
+        if labels[i] != pad_token:
+            last_not_pad_id = i
+            break
+    logits = logits[:last_not_pad_id + 1]
+    labels = labels[:last_not_pad_id + 1]
+    log_probas = scipy.special.log_softmax(logits, axis=1).astype(np.float32)
+    log_probas = [log_probas[i][labels[i]] for i in range(len(labels))]
+    l = np.mean(log_probas)
+    return 2 ** (-l)
+
+
+def perplexity(logits, labels, pad_token=3):
+    pp = []
+    if isinstance(logits, torch.Tensor):
+        logits = logits.detach().cpu().numpy()
+    if isinstance(labels, torch.Tensor):
+        labels = labels.detach().cpu().numpy()
+    for cur_logits, cur_labels in zip(logits, labels):
+        pp.append(_perplexity(np.array(cur_logits), np.array(cur_labels).astype(int), pad_token))
+    return np.mean(pp)