Upload 2 files

code/inference.py

@@ -0,0 +1,66 @@
+from transformers import AutoTokenizer
+from optimum.onnxruntime import ORTModelForCausalLM
+import re
+
+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}
+{char_name}:"""
+
+def model_fn(model_dir):
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = ORTModelForCausalLM.from_pretrained(model_dir, provider="CUDAExecutionProvider")
+    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, return_tensors = "pt").to("cuda")
+        if input_ids.input_ids.size(1) > 1500:
+            chats_curled += 2
+            user_input = user_input[chats_curled:]
+        else: break
+    encoded_output = model.generate(
+        **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"{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
+    }

code/requirements.txt

@@ -0,0 +1 @@
+optimum[onnxruntime-gpu]