[Python] End-to-End ONNX Model with Integrated Tokenization and Decoding in PyTorch

Problem Description:

The objective is to create an "End-to-End ONNX Model with Integrated Tokenization and Decoding".

This involves:

Combining all components: Integrating a tokenizer, a language model (e.g., GPT-2), and a decoder into a single ONNX file.

Goal:

The goal is to have a single ONNX file that can take raw text as input and produce human-readable text as output.

Closest I got is this:

from transformers import GPT2Tokenizer, GPT2LMHeadModel
import onnxruntime as rt
from skl2onnx import convert_sklearn
from sklearn.preprocessing import FunctionTransformer
from sklearn.pipeline import Pipeline
from skl2onnx.common.data_types import StringTensorType

from skl2onnx import to_onnx

model_name = "spital/gpt2-small-czech-cs"
file_name = "results/v3/model/spital-gpt2-small-czech-cs.onnx"

def Tokenizer(input_text):
# Loading the tokenizer
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
# Setting the maximum sequence length
tokenizer.model_max_length = 1024
# PyTorch tensor text tokenization
output = tokenizer.encode(input_text, return_tensors='pt')
return output

tokenizer_transformer = FunctionTransformer(Tokenizer)

# Testing the input tokenizer
input_text = "Téma: Umělá inteligence v moderní společnosti."
tokenized_output = tokenizer_transformer.transform(input_text)
print(tokenized_output)

def Model(tokenized_input):
# Loading the tokenizer
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
# Loading the model
model = GPT2LMHeadModel.from_pretrained(model_name, pad_token_id=tokenizer.eos_token_id)
# Prediction
output = model.generate(tokenized_input, max_length=150, num_return_sequences=1, no_repeat_ngram_size=2, early_stopping=True)
return output

model_transformer = FunctionTransformer(Model)

# Testing the model prediction
predicted_output = model_transformer.transform(tokenized_output)
print(predicted_output)

def Decoder(predicted_input):
# Loading the tokenizer
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
# Setting the maximum sequence length
tokenizer.model_max_length = 1024
# Decoding tokens into readable text
output = tokenizer.decode(predicted_input[0], skip_special_tokens=True)
return output

decoder_transformer = FunctionTransformer(Decoder)

# Decoding testing
decoded_output = decoder_transformer.transform(predicted_output)
print(decoded_output)

pipeline = Pipeline([
('tokenizer', tokenizer_transformer),
('model', model_transformer),
('decoder', decoder_transformer)
])

# Pipeline testing
input_text = "Téma: Umělá inteligence v moderní společnosti."
output_text = pipeline.transform(input_text)
print(output_text)

# Defining input and output shapes
initial_type = [("input", StringTensorType([1]))]
final_type = [("output", StringTensorType([1]))]

# Pipeline conversion to ONNX format
onnx_model = convert_sklearn(pipeline, initial_types=initial_type, final_types=final_type, target_opset=18, options={'zipmap': False})


The problem is that sklearn.preprocessing.FunctionTransformer doesn't support exporting ONNX models from PyTorch. It just supports running PyTorch transformators as python script as you can see and test.

Continue reading...