Merge branch 'bloom-personachat' of https://huggingface.co/spaces/hivemind-personalized-chat/chat-gradio

personalized-chat-bot/bot_example.py

@@ -0,0 +1,60 @@
+import transformers
+
+import argparse
+import json
+
+from petals.client.remote_model import DistributedBloomForCausalLM
+
+from personalized_chat_bot import PersonalizedChatBot, PersonalityManager
+from models.personality_clustering import PersonalityClustering
+
+def load_config(path):
+    with open(path, 'r') as f:
+        config = json.load(f)
+    return argparse.Namespace(**config)
+
+
+def main():
+    greating = 'Describe the person you want to talk:'
+    print(greating)
+    persona_description = input()
+    print('Cool! wait a few seconds...')
+    personality_clustering = PersonalityClustering()
+    personality_clustering.load('./data/models/personality_clustering_500_paraphrase-MiniLM-L6-v2_k-means.pkl')
+
+    hook = lambda dct: {int(k): v for k, v in dct.items()}
+    with open('prompt_paths.json', 'r') as f:
+        prompt_paths = json.load(f, object_hook=hook)
+
+    pm = PersonalityManager(prompt_paths, personality_clustering)
+    prompt_path, closest_persona = pm.get_prompt(persona_description)
+    print(f'The closest personality is: {closest_persona}')
+    print('Wait a little longer...')
+    config = load_config('./scripts/config_176b.json')
+
+    model = DistributedBloomForCausalLM.from_pretrained(
+        config.MODEL_NAME,
+        pre_seq_len=config.NUM_PREFIX_TOKENS,
+        tuning_mode=config.TUNING_MODE
+    ).to(config.DEVICE)
+
+    generation_config = load_config('generation_config.json')
+
+    tokenizer = transformers.BloomTokenizerFast.from_pretrained(config.MODEL_NAME)
+    tokenizer.padding_side = 'right'
+    tokenizer.model_max_length = config.MODEL_MAX_LENGTH
+
+    chatbot = PersonalizedChatBot(model, tokenizer, generation_config=generation_config)
+    chatbot.load_prompt(prompt_path)
+    print('Done! You can start a dialogue.')
+    try:
+        while True:
+            text = input('You: ')
+            answer = chatbot.answer(text)
+            print(f'Bloom: {answer}')
+    except KeyboardInterrupt:
+        print('Thank you for the conversation!')
+
+
+if __name__ == '__main__':
+    main()

personalized-chat-bot/data.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d73016d5eccc0eeb641f623789e6a80c601572aee825603bdfacf84c9e8f705
+size 12635714

personalized-chat-bot/generation_config.json

@@ -0,0 +1 @@
+{"MAX_TOKENS": 16, "TOP_K": 100, "TEMPERATURE": 0.8}

personalized-chat-bot/models/__init__.py

@@ -0,0 +1 @@
+# coding=utf-8

personalized-chat-bot/models/personality_clustering.py

@@ -0,0 +1,74 @@
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from sklearn.cluster import KMeans
+import pickle
+
+
+class PersonalityClustering:
+    DEFAULT_SENTENCE_TRANSFORMER = 'paraphrase-MiniLM-L6-v2'
+
+    @property
+    def sentence_transformer(self):
+        """Ленивая инициализация sentence_transformer."""
+        if not self.__sentence_transformer:
+            self.__sentence_transformer = SentenceTransformer(self.model_name,  device=self.device)
+        return self.__sentence_transformer
+
+    @property
+    def clustering(self):
+        """Ленивая инициализация кластеризации."""
+        if not self.__clustering:
+            self.__clustering = KMeans(n_clusters=self.n_clusters)
+        return self.__clustering
+
+    def __init__(self, n_clusters=None, device='cpu', model_name=None):
+        if model_name is None:
+            self.model_name = self.DEFAULT_SENTENCE_TRANSFORMER
+        else:
+            self.model_name = model_name
+        self.device = device
+        self.n_clusters = n_clusters
+        self._cluster_centers = None
+        self.__clustering = None
+        self.__sentence_transformer = None
+
+    def load(self, path):
+        with open(path, "rb") as f:
+            self.__clustering, self._cluster_centers = pickle.load(f)
+
+    def save(self, path):
+        with open(path, "wb") as f:
+            pickle.dump((self.__clustering, self._cluster_centers), f)
+
+    def fit(self, personalities):
+        personalities = np.array(list(personalities))
+        train_embeddings = self.sentence_transformer.encode(personalities)
+        clusters = self.clustering.fit_predict(train_embeddings)
+        persona_cluster_centers = []
+        for clust, center in enumerate(self.clustering.cluster_centers_):
+            cur_clust_embed = train_embeddings[clusters == clust]
+            cur_clust_personalities = personalities[clusters == clust]
+            min_distance_to_center = np.inf
+            persona_center = None
+            for embed, persona in zip(cur_clust_embed, cur_clust_personalities):
+                cur_distance_to_center = np.linalg.norm(embed - center)
+                if cur_distance_to_center < min_distance_to_center:
+                    min_distance_to_center = cur_distance_to_center
+                    persona_center = persona
+            persona_cluster_centers.append(persona_center)
+        self._cluster_centers = np.array(persona_cluster_centers)
+        return self
+
+    def predict(self, personalities):
+        personalities = np.array(list(personalities))
+        embeddings = self.sentence_transformer.encode(personalities)
+        clusters = self.clustering.predict(embeddings)
+        return clusters
+
+    def predict_nearest_personality(self, personalities):
+        clusters = self.predict(personalities)
+        return np.array([self._cluster_centers[clust] for clust in clusters])
+
+    def fit_predict(self, personalities):
+        self.fit(personalities)
+        return self.predict(personalities)

personalized-chat-bot/personalized_chat_bot.py

@@ -0,0 +1,65 @@
+import argparse
+import json
+import torch
+from sklearn.neighbors import KDTree
+
+
+class PersonalityManager:
+    def __init__(self, prompt_paths, personality_clustering):
+        self.prompt_paths = prompt_paths
+        self.personality_clustering = personality_clustering
+
+        self.persona_ids = list(prompt_paths.keys())
+        self.personalities = [personality_clustering._cluster_centers[i]
+                              for i in self.persona_ids]
+
+        self.embeddings = personality_clustering.sentence_transformer.encode(self.personalities)
+        self._nearest_neighbours = KDTree(self.embeddings, metric='euclidean')
+
+    def get_prompt(self, description):
+        embedding = self.personality_clustering.sentence_transformer.encode([description])
+        dist, ind = self._nearest_neighbours.query(embedding, k=1)
+        persona_id = self.persona_ids[ind[0][0]]
+        prompt_path = self.prompt_paths[persona_id]
+        cluster_center = self.personality_clustering._cluster_centers[persona_id]
+        return prompt_path, cluster_center
+
+
+class PersonalizedChatBot:
+    def __init__(self, model, tokenizer, prompt_path=None, generation_config=None):
+        self.model = model
+        if prompt_path is not None:
+            self.load_prompt(prompt_path)
+        self.tokenizer = tokenizer
+        self.separator = '\n'
+        self.dialog = ''
+        self.generation_config = generation_config
+
+    def load_prompt(self, path):
+        self.model.transformer.prompt_embeddings.load_state_dict(torch.load(path))
+
+    def load_config(self, path):
+        with open(path, 'r') as f:
+            config = json.load(f)
+        self.generation_config = argparse.Namespace(**config)
+
+    def reset_dialog(self, ):
+        self.dialog = ''
+
+    def answer(self, phrase):
+        if len(phrase) == 0:
+            return
+        self.dialog += f"{phrase}{self.separator}"
+        inputs = self.tokenizer([self.dialog], return_tensors='pt')['input_ids']
+        outputs = self.model.generate(
+            inputs,
+            temperature=self.generation_config.TEMPERATURE,
+            do_sample=True,
+            top_k=self.generation_config.TOP_K,
+            eos_token_id=self.tokenizer.eos_token_id,
+            max_new_tokens=self.generation_config.MAX_TOKENS,
+        )
+        bloom_answer = self.tokenizer.batch_decode(outputs)[0]
+        bloom_answer = bloom_answer[len(self.dialog):].split("\n")[0]
+        self.dialog += f"{bloom_answer}{self.separator}"
+        return bloom_answer

personalized-chat-bot/prompt_paths.json

@@ -0,0 +1,16 @@
+{
+  "113": "./data/models/176b/113_persona_prompt_embedding.pt",
+  "54": "./data/models/176b/54_persona_prompt_embedding.pt",
+  "169": "./data/models/176b/169_persona_prompt_embedding.pt",
+  "364": "./data/models/176b/364_persona_prompt_embedding.pt",
+  "214": "./data/models/176b/214_persona_prompt_embedding.pt",
+  "125": "./data/models/176b/125_persona_prompt_embedding.pt",
+  "103": "./data/models/176b/103_persona_prompt_embedding.pt",
+  "200": "./data/models/176b/200_persona_prompt_embedding.pt",
+  "296": "./data/models/176b/296_persona_prompt_embedding.pt",
+  "20": "./data/models/176b/20_persona_prompt_embedding.pt",
+  "384": "./data/models/176b/384_persona_prompt_embedding.pt",
+  "365": "./data/models/176b/365_persona_prompt_embedding.pt",
+  "451": "./data/models/176b/451_persona_prompt_embedding.pt",
+  "80": "./data/models/176b/80_persona_prompt_embedding.pt"
+}

personalized-chat-bot/scripts/__init__.py

@@ -0,0 +1 @@
+# coding=utf-8

personalized-chat-bot/scripts/config_176b.json

@@ -0,0 +1,16 @@
+{
+  "PERSONACHAT_DATASET_NAME": "bavard/personachat_truecased",
+  "MODEL_NAME": "bigscience/bloom-petals",
+  "INITIAL_PEERS": [],
+  "NUM_PREFIX_TOKENS": 16,
+  "DEVICE": "cpu",
+  "BATCH_SIZE": 4,
+  "LR": 0.01,
+  "WEIGHT_DECAY": 0.0,
+  "NUM_SAMPLES": 1000,
+  "SEED": 42,
+  "MODEL_MAX_LENGTH": 256,
+  "TUNING_MODE": "ptune",
+  "N_EPOCH": 10,
+  "PADDING_SIDE": "right"
+}

personalized-chat-bot/scripts/config_6b.json

@@ -0,0 +1,16 @@
+{
+  "PERSONACHAT_DATASET_NAME": "bavard/personachat_truecased",
+  "MODEL_NAME": "bigscience/test-bloomd-6b3",
+  "INITIAL_PEERS":["/ip4/193.106.95.184/tcp/31000/p2p/QmSg7izCDtowVTACbUmWvEiQZNY4wgCQ9T9Doo66K59X6q"],
+  "NUM_PREFIX_TOKENS": 16,
+  "DEVICE": "cpu",
+  "BATCH_SIZE": 4,
+  "LR": 0.01,
+  "WEIGHT_DECAY": 0.0,
+  "NUM_SAMPLES": 1000,
+  "SEED": 42,
+  "MODEL_MAX_LENGTH": 256,
+  "TUNING_MODE": "ptune",
+  "N_EPOCH": 1,
+  "PADDING_SIDE": "right"
+}

personalized-chat-bot/scripts/fit_personality_clustering.py

@@ -0,0 +1,52 @@
+import argparse
+from datasets import load_dataset
+from models.personality_clustering import PersonalityClustering
+import os
+
+"""Пример запуска
+python -m  scripts.fit_personality_clustering --clustering-path data/models --n-clusters 500
+"""
+
+PERSONACHAT_DATASET = "bavard/personachat_truecased"
+
+
+def load_persona_chat_personalities(personachat_dataset):
+    dataset = load_dataset(personachat_dataset)
+    train_personalities = [sent for persona in dataset['train']['personality']
+                           for sent in persona]
+    test_personalities = [sent for persona in dataset['train']['personality']
+                          for sent in persona]
+    personalities = list(set(train_personalities) | set(test_personalities))
+    return personalities
+
+
+def parse_args(args=None):
+    parser = argparse.ArgumentParser(add_help=True, description="Class for personality clustering.")
+
+    parser.add_argument('-clustering-path', '--clustering-path', type=str,
+                        help='Path to clustering data.')
+    parser.add_argument('-n-clusters', '--n-clusters', type=int, default=500,
+                        help='The number of clusters to form.')
+    parser.add_argument('-model-name', '--model-name', type=str, default=None, required=False)
+    args = parser.parse_args(args)
+    return args
+
+
+def main():
+    args = parse_args()
+    personalities = load_persona_chat_personalities(PERSONACHAT_DATASET)
+    print('Data loaded')
+    model = PersonalityClustering(n_clusters=args.n_clusters)
+    print('Model fitting')
+    model.fit(personalities)
+    print('Model fitted')
+    if args.model_name is None:
+        model_name = f'personality_clustering_{model.n_clusters}_{model.model_name}_k-means.pkl'
+    else:
+        model_name = args.model_name
+    model.save(os.path.join(args.clustering_path, model_name))
+    print(f'{model_name} saved')
+
+
+if __name__ == '__main__':
+    main()

personalized-chat-bot/scripts/train_all.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+#python -m scripts.train_bloom_personachat --persona-ids 113 54 169 --config scripts/config_176b.json --prompt-path data/models/176b/ --wandb-project bloom_personachat_176b
+#python -m scripts.train_bloom_personachat --persona-ids 364 214 125 --config scripts/config_176b.json --prompt-path data/models/176b/ --wandb-project bloom_personachat_176b
+#python -m scripts.train_bloom_personachat --persona-ids 103 200 296 --config scripts/config_176b.json --prompt-path data/models/176b/ --wandb-project bloom_personachat_176b
+#python -m scripts.train_bloom_personachat --persona-ids 20 384 365 --config scripts/config_176b.json --prompt-path data/models/176b/ --wandb-project bloom_personachat_176b
+#python -m scripts.train_bloom_personachat --persona-ids 208 43 99 --config scripts/config_176b.json --prompt-path data/models/176b/ --wandb-project bloom_personachat_176b
+#python -m scripts.train_bloom_personachat --persona-ids 426 477 470 --config scripts/config_176b.json --prompt-path data/models/176b/ --wandb-project bloom_personachat_176b
+python -m scripts.train_bloom_personachat --persona-ids 470 --config scripts/config_176b.json --prompt-path data/models/176b/ --wandb-project bloom_personachat_176b
+
+python -m scripts.train_bloom_personachat --persona-ids 329 402 382 --config scripts/config_176b.json --prompt-path data/models/176b/ --wandb-project bloom_personachat_176b

personalized-chat-bot/scripts/train_bloom_personachat.py

@@ -0,0 +1,123 @@
+import argparse
+
+import torch.cuda
+from datasets import load_dataset
+import json
+import os
+import transformers
+from torch.utils.data import Subset
+import wandb
+import numpy as np
+import gc
+
+from models.personality_clustering import PersonalityClustering
+from util.bloom_trainer import BloomTrainer
+from util.data import PersonaChatDataset
+from util.metrics import perplexity
+
+from petals.client.remote_model import DistributedBloomForCausalLM
+
+"""Пример запуска
+python -m scripts.train_bloom_personachat --persona-ids 6 --config scripts/config.json --prompt-path data/models/
+"""
+
+DEFAULT_CLUSTERING_MODEL = './data/models/personality_clustering_500_paraphrase-MiniLM-L6-v2_k-means.pkl'
+MAX_VAL_DATA_SIZE = 4
+
+
+def load_config(path):
+    with open(path, 'r') as f:
+        config = json.load(f)
+    return argparse.Namespace(**config)
+
+
+def main():
+    args = parse_args()
+    persona_clustering = PersonalityClustering()
+    persona_clustering.load(args.clustering_model_path)
+
+    config = load_config(args.config)
+
+    tokenizer = transformers.BloomTokenizerFast.from_pretrained(config.MODEL_NAME)
+    tokenizer.padding_side = config.PADDING_SIDE
+    tokenizer.model_max_length = config.MODEL_MAX_LENGTH
+
+    dataset = load_dataset(config.PERSONACHAT_DATASET_NAME)
+    personachat_train_dataset = PersonaChatDataset(persona_clustering,
+                                                   dataset['train'],
+                                                   tokenizer)
+    personachat_val_dataset = PersonaChatDataset(persona_clustering,
+                                                 dataset['validation'],
+                                                 tokenizer)
+
+    for id in args.persona_ids:
+        prompt_path = os.path.join(args.prompt_path, f'{id}_persona_prompt_embedding.pt')
+        train_dataset = personachat_train_dataset[id]
+        val_dataset = personachat_val_dataset[id]
+        honest_validation = True
+        if len(val_dataset) < 4:
+            val_dataset = personachat_train_dataset[id]
+            honest_validation = False
+        # для ускорения обрежем размер валидации до некоторой границы
+        if len(val_dataset) > MAX_VAL_DATA_SIZE:
+            subset_indexes = np.random.choice(len(val_dataset), MAX_VAL_DATA_SIZE, replace=False)
+            val_dataset = Subset(val_dataset, subset_indexes)
+        # train_dataset.shuffle()
+
+        wandb_run = wandb.init(
+            project=args.wandb_project,
+            config={
+                'lr': config.LR,
+                'batch_size': config.BATCH_SIZE,
+                'persona_id': id,
+                'device': config.DEVICE,
+                'model_name': config.MODEL_NAME,
+                'n_epoch': config.N_EPOCH,
+                'honest_validation': honest_validation
+            },
+            name=f'id{id}',
+            reinit=True
+        )
+        if len(config.INITIAL_PEERS) == 0:
+            model = DistributedBloomForCausalLM.from_pretrained(
+                config.MODEL_NAME,
+                pre_seq_len=config.NUM_PREFIX_TOKENS,
+                tuning_mode=config.TUNING_MODE
+            ).to(config.DEVICE)
+        else:
+            model = DistributedBloomForCausalLM.from_pretrained(
+                config.MODEL_NAME,
+                initial_peers=config.INITIAL_PEERS,
+                pre_seq_len=config.NUM_PREFIX_TOKENS,
+                tuning_mode=config.TUNING_MODE
+            ).to(config.DEVICE)
+
+        trainer = BloomTrainer(model, config, train_dataset, val_dataset, wandb_run, prompt_path)
+        trainer.train()
+        eval_perplexity = trainer.evaluate(perplexity)
+        trainer.save_model(prompt_path)
+        wandb_run.log({'perplexity': eval_perplexity, 'model_path': prompt_path})
+
+        del model
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+def parse_args(args=None):
+    parser = argparse.ArgumentParser(add_help=True,
+                                     description="bloom training script")
+    parser.add_argument('--persona-ids', type=int, nargs='+',
+                        help='Ids of persona')
+    parser.add_argument('-clustering-model-path', '--clustering-model-path', type=str,
+                        default=DEFAULT_CLUSTERING_MODEL,
+                        help='Path to clustering model')
+    parser.add_argument('--config', type=str, help='Path to training config file')
+    parser.add_argument('--prompt-path', type=str,
+                        help='Path to dir with trained soft prompts')
+    parser.add_argument('--wandb-project', type=str, default='test_bloom_personachat_176b_v3')
+    args = parser.parse_args(args)
+    return args
+
+
+if __name__ == '__main__':
+    main()

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)