Spaces:
Runtime error
Runtime error
File size: 8,612 Bytes
e71a2ba dd71c9a 91ffd2c 776e43c e71a2ba 776e43c 0669a02 aae4195 dd71c9a aae4195 776e43c 91ffd2c dd71c9a aae4195 776e43c b649ec8 aae4195 b649ec8 a4f8b32 b649ec8 a4f8b32 b649ec8 a4f8b32 b649ec8 a4f8b32 b649ec8 91ffd2c dd71c9a 91ffd2c dd71c9a 91ffd2c dd71c9a 91ffd2c dd71c9a 91ffd2c dd71c9a 91ffd2c dd71c9a 91ffd2c dd71c9a b649ec8 aae4195 b649ec8 aae4195 b649ec8 aae4195 b649ec8 aae4195 b649ec8 aae4195 b649ec8 aae4195 b649ec8 776e43c b649ec8 dd71c9a aae4195 b649ec8 aae4195 b649ec8 aae4195 0669a02 aae4195 776e43c a4f8b32 aae4195 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 |
import sys
import json
import argparse
import gradio as gr
import torch
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 personalized_chat_bot import PersonalizedChatBot, PersonalityManager
from models.personality_clustering import PersonalityClustering
MODEL_NAME = "bigscience/bloom-petals"
tokenizer_bloomd = transformers.BloomTokenizerFast.from_pretrained(MODEL_NAME)
model_bloomd = DistributedBloomForCausalLM.from_pretrained(MODEL_NAME, low_cpu_mem_usage=True)
tokenizer_DialoGPT_small = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
model_DialoGPT_small = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
tokenizer_DialoGPT_medium = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model_DialoGPT_medium = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
tokenizer_DialoGPT_large = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
model_DialoGPT_large = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
def predict_common_bloom(model, tokenizer, input_text, history, person_description, number_of_new_tokens):
new_user_input_ids = tokenizer.encode(input_text + '\n', return_tensors='pt')
person_description_ids = tokenizer.encode(person_description + '\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}')
input_with_desc_ids = torch.cat([person_description_ids, bot_input_ids], dim=-1)
history = model.generate(
input_with_desc_ids,
max_new_tokens=number_of_new_tokens,
pad_token_id=tokenizer.eos_token_id
).tolist()
print(f'history: {history}')
history[0] = history[0][len(person_description_ids[0]):]
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 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):
personality_clustering = PersonalityClustering()
personality_clustering.load('personalized-chat-bot/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('personalized-chat-bot/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(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')
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)
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}')
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):
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')
bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)
input_with_desc_ids = torch.cat([person_description_ids, bot_input_ids], dim=-1)
history = model.generate(
input_with_desc_ids,
max_new_tokens=number_of_new_tokens,
pad_token_id=tokenizer.eos_token_id
).tolist()
history[0] = history[0][len(person_description_ids[0]):]
response = tokenizer.decode(history[0]).split("<|endoftext|>")
response = [(response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)]
return response, history
def predict(
input_text,
history=None,
person_description=None,
number_of_new_tokens=10,
model_name=None,
del_hist=None
):
if history is None or del_hist == 'delete history':
history = []
if model_name == 'DialoGPT-small':
model = model_DialoGPT_small
tokenizer = tokenizer_DialoGPT_small
return predict_dialo_gpt(model, tokenizer, input_text, history, person_description, number_of_new_tokens)
elif model_name == 'DialoGPT-medium':
model = model_DialoGPT_medium
tokenizer = tokenizer_DialoGPT_medium
return predict_dialo_gpt(model, tokenizer, input_text, history, person_description, number_of_new_tokens)
elif model_name == 'DialoGPT-large':
model = model_DialoGPT_large
tokenizer = tokenizer_DialoGPT_large
return predict_dialo_gpt(model, tokenizer, input_text, history, person_description, number_of_new_tokens)
elif model_name == 'bloom-petals':
model = model_bloomd
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
tokenizer = tokenizer_DialoGPT_medium
return predict_dialo_gpt(model, tokenizer, input_text, history, person_description, number_of_new_tokens)
gr.Interface(
fn=predict,
inputs=[
gr.Textbox(label='Input message', lines=1, placeholder="Enter your message..."),
"state",
gr.Textbox(label='Person Description', lines=2, placeholder="Enter a description of the person..."),
gr.Slider(label='Number of new tokens', minimum=2, maximum=100, value=10),
gr.Radio(
label='Model name',
choices=[
'DialoGPT-small',
'DialoGPT-medium',
'DialoGPT-large',
'bloom-petals',
'bloom-petals-cluster',
]
),
gr.Radio(
label='Delete history',
value="Don't delete history",
choices=[
'delete history',
"Don't delete history"
]),
],
outputs=[gr.Chatbot(label='History of the dialogue'), "state"],
).launch(),
|