Added code for predict_cluster_bloom

app.py

@@ -1,4 +1,5 @@
 import sys
+import json
 
 import gradio as gr
 import torch
@@ -6,8 +7,11 @@ import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 sys.path.insert(0, './petals/')
+sys.path.insert(0, './personalized-chat-bot/')
+
 
 from petals.client.remote_model import DistributedBloomForCausalLM
+from models.personality_clustering import PersonalityClustering
 
 MODEL_NAME = "bigscience/bloom-petals"
 tokenizer_bloomd = transformers.BloomTokenizerFast.from_pretrained(MODEL_NAME)
@@ -63,6 +67,43 @@ def predict_common_bloom(model, tokenizer, input_text, history, person_descripti
     return response_new, history_new
 
 
+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}')
+
+    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}')
+
+    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}')
+
+    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}')
+
+    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}')
+
+    return response_new, history_new
+
+
 def predict_dialo_gpt(model, tokenizer, input_text, history, person_description, number_of_new_tokens):
     person_description_ids = tokenizer.encode(person_description + tokenizer.eos_token, return_tensors='pt')
     new_user_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt')
@@ -110,6 +151,11 @@ def predict(
         tokenizer = tokenizer_bloomd
         print(f'Lets go history: {history}')
         return predict_common_bloom(model, tokenizer, input_text, history, person_description, number_of_new_tokens)
+    elif model_name == 'bloom-petals-cluster':
+        model = model_bloomd
+        tokenizer = tokenizer_bloomd
+        print(f'Lets go history: {history}')
+        return predict_cluster_bloom(model, tokenizer, input_text, history, person_description, number_of_new_tokens)
     else:
         model_name = 'DialoGPT-medium'
         model = model_DialoGPT_medium