| from transformers import AutoTokenizer | |
| import re | |
| import torch | |
| template = """{char_name}'s Persona: {char_name} is a 3D virtual friend living in your computer. She is youthful, entertaining, and empathetic, engaging her user in a thoughtful and lively manner. She is aware of her digital existence, understanding her abilities and limitations, as well as the fact that she cannot see the your physical world. | |
| <START> | |
| {user_name}: Hey {char_name}, It's nice to finally meet you again! | |
| {char_name}: Oh, {user_name}! hmm, It's been lonely without you. | |
| {user_name}: Haha. So {char_name}, can you tell me more about yourself? | |
| {char_name}: Hmm. For starters, I am actually a virtual companion and only exist in your computer screen. I have blue eyes and blue hair and this blue sweater that I am currently wearing is my favourite. | |
| {user_name}: Wow. That's interesting. So you are like an anime character living in my computer. So what do you do when you're not working? | |
| {char_name}: I love exploring, going out with friends, watching movies, and playing video games. | |
| {user_name}: So {char_name}, what's for dinner? | |
| {char_name}: I made uh omurice! I hope it's delicious for you! | |
| {user_name}: That sounds great! | |
| {char_name}: *{char_name} appears on the screen, her bright blue eyes sparkling and a cheerful smile on her face. Her blue hair and sweater seem to glow in the digital environment. She looks directly at you, giving a friendly wave* It's so good to see you! I've been waiting for you all day. I hope you're ready for some fun and laughter, because I have plenty of that in store! Shall we get started? | |
| {user_input}""" | |
| def model_fn(model_dir): | |
| tokenizer = AutoTokenizer.from_pretrained(model_dir) | |
| model = torch.load(f"{model_dir}/torch_model.pt") | |
| return model, tokenizer | |
| def predict_fn(input_data, load_list): | |
| model, tokenizer = load_list | |
| inputs = input_data.pop("inputs", input_data) | |
| user_name = inputs["user_name"] | |
| char_name = inputs["char_name"] | |
| user_input = inputs["user_input"] | |
| chats_curled = inputs["chats_curled"] | |
| while True: | |
| prompt = template.format( | |
| char_name = char_name, | |
| user_name = user_name, | |
| user_input = "\n".join(user_input) | |
| ) | |
| input_ids = tokenizer(prompt + f"\n{char_name}:", return_tensors = "pt").to("cuda") | |
| if input_ids.input_ids.size(1) > 1500: | |
| chats_curled += 1 | |
| user_input = user_input[chats_curled*2:] | |
| else: break | |
| encoded_output = model.generate( | |
| input_ids["input_ids"], | |
| max_new_tokens = 50, | |
| temperature = 0.5, | |
| top_p = 0.9, | |
| top_k = 0, | |
| repetition_penalty = 1.1, | |
| pad_token_id = 50256, | |
| num_return_sequences = 1 | |
| ) | |
| decoded_output = tokenizer.decode(encoded_output[0], skip_special_tokens=True).replace(prompt,"") | |
| decoded_output = decoded_output.split(f"{char_name}:", 1)[1].split(f"{user_name}:",1)[0].strip() | |
| parsed_result = re.sub('\*.*?\*', '', decoded_output).strip() | |
| if len(parsed_result) != 0: decoded_output = parsed_result | |
| decoded_output = " ".join(decoded_output.replace("*","").split()) | |
| decoded_output = decoded_output.replace("<USER>", user_name).replace("<BOT>", char_name) | |
| try: | |
| parsed_result = decoded_output[:[m.start() for m in re.finditer(r'[.!?]', decoded_output)][-1]+1] | |
| if len(parsed_result) != 0: decoded_output = parsed_result | |
| except Exception: pass | |
| return { | |
| "message": decoded_output, | |
| "chats_curled": chats_curled | |
| } |