Added Cluster Bloom

app.py

@@ -1,5 +1,6 @@
 import sys
 import json
+import argparse
 
 import gradio as gr
 import torch
@@ -11,6 +12,7 @@ sys.path.insert(0, './personalized-chat-bot/')
 
 
 from petals.client.remote_model import DistributedBloomForCausalLM
+from personalized_chat_bot import PersonalizedChatBot, PersonalityManager
 from models.personality_clustering import PersonalityClustering
 
 MODEL_NAME = "bigscience/bloom-petals"
@@ -67,41 +69,59 @@ def predict_common_bloom(model, tokenizer, input_text, history, person_descripti
     return response_new, history_new
 
 
+def load_config(path):
+    with open(path, 'r') as f:
+        config = json.load(f)
+    return argparse.Namespace(**config)
+
+
 def predict_cluster_bloom(model, tokenizer, input_text, history, person_description, number_of_new_tokens):
-    new_user_input_ids = tokenizer.encode(input_text + '\n', return_tensors='pt')
-    print('Started predict_common_bloom')
-    print(f'history: {history}')
-    if history != []:
-        bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)
-    else:
-        bot_input_ids = new_user_input_ids
-    print(f'bot_input_ids: {bot_input_ids}')
+    personality_clustering = PersonalityClustering()
+    personality_clustering.load('personalized-chat-bot/data/models/personality_clustering_500_paraphrase-MiniLM-L6-v2_k-means.pkl')
 
-    history = model.generate(
-        bot_input_ids,
-        max_new_tokens=number_of_new_tokens,
-        pad_token_id=tokenizer.eos_token_id
-    ).tolist()
-    print(f'history: {history}')
+    hook = lambda dct: {int(k): v for k, v in dct.items()}
+    with open('personalized-chat-bot/prompt_paths.json', 'r') as f:
+        prompt_paths = json.load(f, object_hook=hook)
 
-    decode_all = tokenizer.decode(history[0][:len(bot_input_ids[0])])
-    all_responses = tokenizer.decode(history[0][len(bot_input_ids[0]):]).split('\n')
-    if all_responses[0]:
-        decode_all += all_responses[0] + '\n'
-    else:
-        decode_all += all_responses[1] + '\n'
-    print(f'decode_all: {decode_all}')
+    pm = PersonalityManager(prompt_paths, personality_clustering)
+    prompt_path, closest_persona = pm.get_prompt(person_description)
+    print(f'The closest personality is: {closest_persona}')
+    print('Wait a little longer...')
+    config = load_config('personalized-chat-bot/scripts/config_176b.json')
 
-    history_new = tokenizer.encode(decode_all, return_tensors='pt')
-    print(f'history_new: {history_new}')
 
-    decode_all_split = decode_all.split('\n')
-    print(f'decode_all_split: {decode_all_split}')
+    model = DistributedBloomForCausalLM.from_pretrained(
+        config.MODEL_NAME,
+        pre_seq_len=config.NUM_PREFIX_TOKENS,
+        tuning_mode=config.TUNING_MODE,
+        # max_new_tokens=number_of_new_tokens,
+    ).to(config.DEVICE)
 
-    response_new = [(decode_all_split[i], decode_all_split[i + 1]) for i in range(0, len(decode_all_split) - 1, 2)]
-    print(f'response_new: {response_new}')
+    generation_config = load_config('personalized-chat-bot/generation_config.json')
+    generation_config.max_new_tokens=number_of_new_tokens
+    print(f'generation_config: {generation_config}')
 
-    return response_new, history_new
+    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('personalized-chat-bot/' + prompt_path)
+    if history != []:
+        input_text = tokenizer.decode(history[0]) + '\n' + input_text
+    print(f'INPUT: {input_text}')
+    output = chatbot.answer(input_text)
+    all_text = input_text + '\n' + output
+    print(f'all_text: {all_text}')
+
+    history = tokenizer.encode(all_text,  return_tensors='pt')
+    print(f'history: {history}')
+
+    response = tokenizer.decode(history[0]).split("\n")
+    response = [(response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)]
+    print(f'response: {response}')
+
+    return response, history
 
 
 def predict_dialo_gpt(model, tokenizer, input_text, history, person_description, number_of_new_tokens):