<jupyter_start><jupyter_text>Conversion to ONNXONNX is a different format for running machine learning models. The ONNX format is much faster on CPU, sometimes 5 times as fast as PyTorch!While the EAWSW model is designed to be small, accurate and accessible, for some people it's still too much to run...Hosting the model as a free service for players is an option. An ONNX version of the model allows us to host the model on CPU yet have faster response times! Given that the model is made in a time with chip shortage, running on hardware I already have inside a server is efficient, scalable and cheaper.An important note is that ONNX doesn't execute logic by itself, and you have to do that yourself, `onnx_model_manager.py` intends to deal with this for us.<jupyter_code>%load_ext autoreload
%autoreload 2

from model_utils import train_model, split_data, split_branches, get_model, set_pretrained_model_dropout, get_dataset, ModelSeeder
from config import Config
import json
import matplotlib.pyplot as plt
%matplotlib inline
import math
import random
import time
import onnx
import logging
from onnx_model_manager import OnnxModelManager
from onnxruntime.quantization import quantize_dynamic, QuantType
import os
import datasets
import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM
from model_manager import ModelManager
saved_model_path = os.path.join("models", "awsw_mixed")
saved_model_onnx_path = os.path.join("models", "awsw_onnx")
if not os.path.exists(os.path.join(saved_model_path, "special_tokens_map.json")):
    print("Copying config files from huggingface (needed for conversion)... WARNING: this assumes the structure of the model isn't changed!")
    !cd $saved_model_path && git clone https://huggingface.co/EleutherAI/gpt-neo-125M
    !cp -n $saved_model_path/gpt-neo-125M/* $saved_model_path
    !rm -rf $saved_model_path/gpt-neo-125M
if not os.path.exists(os.path.join(saved_model_onnx_path, "model.onnx")):
    !python3 -m transformers.onnx --model=$saved_model_path --feature=causal-lm-with-past $saved_model_onnx_path
def optimize_onnx():
    model_quant = os.path.join(saved_model_onnx_path, "model_quant.onnx")
    if not os.path.exists(model_quant):
        model_fp32 = os.path.join(saved_model_onnx_path, "model.onnx")
        model_opt = os.path.join(saved_model_onnx_path, "model-opt.onnx")
        quantized_model = quantize_dynamic(model_fp32, model_quant, weight_type = QuantType.QInt8)
        #!rm $model_opt
optimize_onnx()
# Tell pytorch to run this model on the GPU.
device_name = "cuda:0" if torch.cuda.is_available() else "cpu"
device = torch.device(device_name)

onnx_model_manager = OnnxModelManager(os.path.join(saved_model_onnx_path, "model-opt.onnx"))
onnx_model_manager_quant = OnnxModelManager(os.path.join(saved_model_onnx_path, "model_quant.onnx"))
tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neo-125M')
model = AutoModelForCausalLM.from_pretrained(saved_model_path)
model.to(device)
model.eval()
model_manager = ModelManager(model=model, tokenizer=tokenizer, device=device)
print(f"Pretrained model loaded on {device_name}")
prompt = "In my dreams, I'm a dragon"
for i in range(2):
    print("ONNX:", onnx_model_manager.say_raw(prompt, do_sample=True))
    print("ONNX (Quantized):", onnx_model_manager_quant.say_raw(prompt, do_sample=True))
    print("PyTorch:", model_manager.say_raw(prompt, 50, 0.7))
    print('-' * 100)<jupyter_output><empty_output><jupyter_text>TestingWe created a few past (for context) + present prompts (player input) and see the different reactions. This way, we can test the models across different iterations.The first test involves a old prompt to compare the pre-trained model with the one trained on AWSW. Did it manage to store it's data well? Is it able to write down things that have nothing to do with AWSW? (So we know we didn't overfit). **This test generates boring and repetetive** replies! It's because we use no good sampling algorithm, but it does give us a indication of what the model has learned!<jupyter_code>prompts = [
    ('<p><msg>c "Hey Remy!"<d><scn>park2<msg>Ry "Hello, [player_name]."', "How are you?"),
    ('<p><msg>c "I was with Lorem today."<d><scn>park2<msg>Ad "Very nice."', "What do you think of Lorem?"),
    ('<p><msg>m "In Tatsu park, Adine and I sat down."', "Oh my god, Adine. What is this?"),
    ('<p><msg>m "I sat down on a chair in Anna\'s lab."', "What will we do here?"),
]

for (past, prompt) in prompts:
    print(f"Prompt: {prompt}")
    reply = model_manager.say(past, prompt)
    print(f"[Pytorch] Reply: {reply}\n")
    reply = onnx_model_manager.say(past, prompt)
    print(f"[ONNX] Reply: {reply}\n")
    reply = onnx_model_manager_quant.say(past, prompt)
    print(f"[ONNX Quantized] Reply: {reply}\n")
    print("-" * 10)<jupyter_output><empty_output><jupyter_text>Sampling testThis is gonna be interesting!<jupyter_code>for (past, prompt) in prompts:
    print(f"Prompt: {prompt}")
    reply = model_manager.say(past, prompt, top_k = 50, top_p = 0.7)
    print(f"[Pytorch] Reply: {reply}\n")
    reply = onnx_model_manager.say(past, prompt, do_sample = True)
    print(f"[ONNX] Reply: {reply}\n")
    reply = onnx_model_manager_quant.say(past, prompt, do_sample = True)
    print(f"[ONNX Quantized] Reply: {reply}\n")
    print("-" * 10)<jupyter_output><empty_output><jupyter_text>RP testTesting out the injected roleplay actions<jupyter_code>test_rps = [
    "Visit Lorem",
    "Meet with Lorem",
    "Visit Adine",
    "Fight Maverick",
    "Fight Adine",
    "Attack Adine"
]

for rp in test_rps:
    print(f'[Pytorch] {rp} -> {model_manager.say("", rp, top_k = 50, top_p = 0.7)}')
    print(f'[ONNX] {rp} -> {onnx_model_manager.say("", rp, do_sample = True)}')
    print(f'[ONNX Quantized] {rp} -> {onnx_model_manager_quant.say("", rp, do_sample = True)}')
    print("-" * 10)
    
print("Lowercase test")

for rp in test_rps:
    rp = rp[0].lower() + rp[1:]
    print(f'[Pytorch] {rp} -> {model_manager.say("", rp, top_k = 50, top_p = 0.7)}')
    print(f'[ONNX] {rp} -> {onnx_model_manager.say("", rp, do_sample = True)}')
    print(f'[ONNX Quantized] {rp} -> {onnx_model_manager_quant.say("", rp, do_sample = True)}')
    rp = rp.lower()
    print(f'[Pytorch] {rp} -> {model_manager.say("", rp, top_k = 50, top_p = 0.7)}')
    print(f'[ONNX] {rp} -> {onnx_model_manager.say("", rp, do_sample = True)}')
    print(f'[ONNX Quantized] {rp} -> {onnx_model_manager_quant.say("", rp, do_sample = True)}')
    print("-" * 10)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12675626145387905, 'fraction_numerical': 0.00870494807574832, 'mean_word_length': 5.0193014705882355, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 40, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '76520', 'license': 'MIT', 'n_tokens_mistral': 2377, 'n_tokens_neox': 2274, 'n_words': 644, 'path': 'Research/Step 3 - Distilling and converting to ONNX.ipynb', 'repo_name': 'peterwilli/Endless-AWSW', 'size': '6548'}