add personalized-chat-bot

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/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()