add personalized-chat-bot

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