cleaner files

.gitignore

@@ -1,2 +1,4 @@
 venv
-.vscode
+.vscode
+s2v_reddit_2015_md.tar.gz
+__pycache__

FastT5/__init__.py

@@ -0,0 +1,4 @@
+from .huggingface_utils import set_auth_token
+from .onnx_models import OnnxT5, export_and_get_onnx_model, get_onnx_model
+from .ort_settings import get_onnx_runtime_sessions
+from .onnx_exporter import generate_onnx_representation, quantize

FastT5/huggingface_utils.py

@@ -0,0 +1,25 @@
+_auth_token = None
+
+def set_auth_token(token):
+    """Set the token which allows the user to authenticate to hugginface.co for downloading private models
+
+    Args:
+        token (Union[str, bool]): The token value to store. One of:
+            - an API key (from https://huggingface.co/organizations/ORGNAME/settings/token),
+            - a login token obtained by running `$ transformers-cli login`
+            - `True`, which tells transformers to use the login token stored in ~/.huggingface/token
+
+    Returns:
+        None
+    """
+    global _auth_token
+    _auth_token = token
+
+def get_auth_token():
+    """Get the user-configurable auth token, which defaults to None
+
+    Returns:
+        auth_token (Optional[Union[str, bool]]) for authenticating with huggingface.co
+    """
+    global _auth_token
+    return _auth_token

FastT5/mcq.py

@@ -0,0 +1,311 @@
+from flashtext import KeywordProcessor
+from nltk.tokenize import sent_tokenize
+from similarity.normalized_levenshtein import NormalizedLevenshtein
+from nltk.corpus import stopwords
+import torch
+from collections import OrderedDict
+import string
+import pke
+import nltk
+import random
+nltk.download('brown')
+nltk.download('stopwords')
+nltk.download('popular')
+
+
+def MCQs_available(word, s2v):
+    word = word.replace(" ", "_")
+    sense = s2v.get_best_sense(word)
+    if sense is not None:
+        return True
+    else:
+        return False
+
+
+def edits(word):
+    "All edits that are one edit away from `word`."
+    letters = 'abcdefghijklmnopqrstuvwxyz '+string.punctuation
+    splits = [(word[:i], word[i:]) for i in range(len(word) + 1)]
+    deletes = [L + R[1:] for L, R in splits if R]
+    transposes = [L + R[1] + R[0] + R[2:] for L, R in splits if len(R) > 1]
+    replaces = [L + c + R[1:] for L, R in splits if R for c in letters]
+    inserts = [L + c + R for L, R in splits for c in letters]
+    return set(deletes + transposes + replaces + inserts)
+
+
+def sense2vec_get_words(word, s2v):
+    output = []
+
+    word_preprocessed = word.translate(
+        word.maketrans("", "", string.punctuation))
+    word_preprocessed = word_preprocessed.lower()
+
+    word_edits = edits(word_preprocessed)
+
+    word = word.replace(" ", "_")
+
+    sense = s2v.get_best_sense(word)
+    most_similar = s2v.most_similar(sense, n=15)
+
+    compare_list = [word_preprocessed]
+    for each_word in most_similar:
+        append_word = each_word[0].split("|")[0].replace("_", " ")
+        append_word = append_word.strip()
+        append_word_processed = append_word.lower()
+        append_word_processed = append_word_processed.translate(
+            append_word_processed.maketrans("", "", string.punctuation))
+        if append_word_processed not in compare_list and word_preprocessed not in append_word_processed and append_word_processed not in word_edits:
+            output.append(append_word.title())
+            compare_list.append(append_word_processed)
+
+    out = list(OrderedDict.fromkeys(output))
+
+    return out
+
+
+def get_options(answer, s2v):
+    distractors = []
+
+    try:
+        distractors = sense2vec_get_words(answer, s2v)
+        if len(distractors) > 0:
+            print(" Sense2vec_distractors successful for word : ", answer)
+            return distractors, "sense2vec"
+    except:
+        print(" Sense2vec_distractors failed for word : ", answer)
+
+    return distractors, "None"
+
+
+def tokenize_sentences(text):
+    sentences = [sent_tokenize(text)]
+    sentences = [y for x in sentences for y in x]
+    # Remove any short sentences less than 20 letters.
+    sentences = [sentence.strip()
+                 for sentence in sentences if len(sentence) > 20]
+    return sentences
+
+
+def get_sentences_for_keyword(keywords, sentences):
+    keyword_processor = KeywordProcessor()
+    keyword_sentences = {}
+    for word in keywords:
+        word = word.strip()
+        keyword_sentences[word] = []
+        keyword_processor.add_keyword(word)
+    for sentence in sentences:
+        keywords_found = keyword_processor.extract_keywords(sentence)
+        for key in keywords_found:
+            keyword_sentences[key].append(sentence)
+
+    for key in keyword_sentences.keys():
+        values = keyword_sentences[key]
+        values = sorted(values, key=len, reverse=True)
+        keyword_sentences[key] = values
+
+    delete_keys = []
+    for k in keyword_sentences.keys():
+        if len(keyword_sentences[k]) == 0:
+            delete_keys.append(k)
+    for del_key in delete_keys:
+        del keyword_sentences[del_key]
+
+    return keyword_sentences
+
+
+def is_far(words_list, currentword, thresh, normalized_levenshtein):
+    threshold = thresh
+    score_list = []
+    for word in words_list:
+        score_list.append(normalized_levenshtein.distance(
+            word.lower(), currentword.lower()))
+    if min(score_list) >= threshold:
+        return True
+    else:
+        return False
+
+
+def filter_phrases(phrase_keys, max, normalized_levenshtein):
+    filtered_phrases = []
+    if len(phrase_keys) > 0:
+        filtered_phrases.append(phrase_keys[0])
+        for ph in phrase_keys[1:]:
+            if is_far(filtered_phrases, ph, 0.7, normalized_levenshtein):
+                filtered_phrases.append(ph)
+            if len(filtered_phrases) >= max:
+                break
+    return filtered_phrases
+
+
+def get_nouns_multipartite(text):
+    out = []
+
+    extractor = pke.unsupervised.MultipartiteRank()
+    extractor.load_document(input=text, language='en')
+    pos = {'PROPN', 'NOUN'}
+    stoplist = list(string.punctuation)
+    stoplist += stopwords.words('english')
+    extractor.candidate_selection(pos=pos)
+    # 4. build the Multipartite graph and rank candidates using random walk,
+    #    alpha controls the weight adjustment mechanism, see TopicRank for
+    #    threshold/method parameters.
+    try:
+        extractor.candidate_weighting(alpha=1.1,
+                                      threshold=0.75,
+                                      method='average')
+    except:
+        return out
+
+    keyphrases = extractor.get_n_best(n=10)
+
+    for key in keyphrases:
+        out.append(key[0])
+
+    return out
+
+
+def get_phrases(doc):
+    phrases = {}
+    for np in doc.noun_chunks:
+        phrase = np.text
+        len_phrase = len(phrase.split())
+        if len_phrase > 1:
+            if phrase not in phrases:
+                phrases[phrase] = 1
+            else:
+                phrases[phrase] = phrases[phrase]+1
+
+    phrase_keys = list(phrases.keys())
+    phrase_keys = sorted(phrase_keys, key=lambda x: len(x), reverse=True)
+    phrase_keys = phrase_keys[:50]
+    return phrase_keys
+
+
+def get_keywords(nlp, text, max_keywords, s2v, fdist, normalized_levenshtein, no_of_sentences):
+    doc = nlp(text)
+    max_keywords = int(max_keywords)
+
+    keywords = get_nouns_multipartite(text)
+    keywords = sorted(keywords, key=lambda x: fdist[x])
+    keywords = filter_phrases(keywords, max_keywords, normalized_levenshtein)
+
+    phrase_keys = get_phrases(doc)
+    filtered_phrases = filter_phrases(
+        phrase_keys, max_keywords, normalized_levenshtein)
+
+    total_phrases = keywords + filtered_phrases
+
+    total_phrases_filtered = filter_phrases(total_phrases, min(
+        max_keywords, 2*no_of_sentences), normalized_levenshtein)
+
+    answers = []
+    for answer in total_phrases_filtered:
+        if answer not in answers and MCQs_available(answer, s2v):
+            answers.append(answer)
+
+    answers = answers[:max_keywords]
+    return answers
+
+
+def generate_questions_mcq(keyword_sent_mapping, device, tokenizer, model, sense2vec, normalized_levenshtein):
+    batch_text = []
+
+    answers = keyword_sent_mapping.keys()
+    for answer in answers:
+        txt = keyword_sent_mapping[answer]
+        txt_str = "\n".join(txt)
+        context = "context: " + txt_str
+        text = context + " " + "answer: " + answer + " </s>"
+        batch_text.append(text)
+    print(batch_text)
+
+    encoding = tokenizer.batch_encode_plus(
+        batch_text, pad_to_max_length=True, return_tensors="pt")
+
+    print("Running model for generation")
+    input_ids, attention_masks = encoding["input_ids"].to(
+        device), encoding["attention_mask"].to(device)
+
+    with torch.no_grad():
+        outs = model.generate(input_ids=input_ids,
+                              attention_mask=attention_masks,
+                              max_length=150)
+
+    output_array = {}
+    output_array["questions"] = []
+#     print(outs)
+    for index, val in enumerate(answers):
+        individual_question = {}
+        out = outs[index, :]
+        dec = tokenizer.decode(out, skip_special_tokens=True,
+                               clean_up_tokenization_spaces=True)
+
+        Question = dec.replace("question:", "")
+        Question = Question.strip()
+        individual_question["question_statement"] = Question
+        individual_question["question_type"] = "MCQ"
+        individual_question["answer"] = val
+        individual_question["id"] = index+1
+        individual_question["options"], individual_question["options_algorithm"] = get_options(
+            val, sense2vec)
+
+        individual_question["options"] = filter_phrases(
+            individual_question["options"], 10, normalized_levenshtein)
+        index = 3
+        individual_question["extra_options"] = individual_question["options"][index:]
+        individual_question["options"] = individual_question["options"][:index]
+        individual_question["context"] = keyword_sent_mapping[val]
+
+        if len(individual_question["options"]) > 0:
+            output_array["questions"].append(individual_question)
+
+    return output_array
+
+
+# for normal one word questions
+def generate_normal_questions(keyword_sent_mapping, device, tokenizer, model):
+    batch_text = ""
+    answers = keyword_sent_mapping.keys()
+    for answer in answers:
+        txt = keyword_sent_mapping[answer]
+        context = "context: " + txt
+        text = context + " " + "answer: " + answer + " </s>"
+        batch_text.join(text)
+
+    encoding = tokenizer.batch_encode_plus(
+        batch_text, pad_to_max_length=True, return_tensors="pt")
+
+    print("Running model for generation")
+    input_ids, attention_masks = encoding["input_ids"].to(
+        device), encoding["attention_mask"].to(device)
+
+    with torch.no_grad():
+        outs = model.generate(input_ids=input_ids,
+                              attention_mask=attention_masks,
+                              max_length=150)
+
+    output_array = {}
+    output_array["questions"] = []
+
+    for index, val in enumerate(answers):
+        individual_quest = {}
+        out = outs[index, :]
+        dec = tokenizer.decode(out, skip_special_tokens=True,
+                               clean_up_tokenization_spaces=True)
+
+        Question = dec.replace('question:', '')
+        Question = Question.strip()
+
+        individual_quest['Question'] = Question
+        individual_quest['Answer'] = val
+        individual_quest["id"] = index+1
+        individual_quest["context"] = keyword_sent_mapping[val]
+
+        output_array["questions"].append(individual_quest)
+
+    return output_array
+
+
+def random_choice():
+    a = random.choice([0, 1])
+    return bool(a)

FastT5/model_testing_tools.py

@@ -0,0 +1,103 @@
+from time import perf_counter as pc
+from matplotlib import pyplot as plt
+from transformers import AutoTokenizer
+
+import numpy as np
+
+
+def speed_test(
+    onnx_model,
+    torch_model,
+    beam_range: range = range(1, 10, 1),
+    seq_length_range: range = range(10, 500, 50),
+    input_text=None,
+):
+    """
+        method prints the time took for onnx and pytorch model to finish a text generation task
+
+    args:
+        input_text (str) : text input for the model.
+        onnx_model : onnx representation of the t5 model,
+        torch_model : torch represention of the t5 model,
+        beam_range (range) : provide a range, which takes starting end and steps (don't start with 0)
+        sequence_length-range (range) : takes the start, end and steps as a range (start with 10)
+    return :
+        onnx_model_latency : numpy array of latency for each beam number and sequence length
+        pytorch_model_latency : numpy array of latency for each beam number and sequence length
+    """
+
+    if input_text is None:
+        input_text = """translate English to French: A nucleus is a collection of a large number of up and down quarks, confined into triplets (neutrons and protons). According to the strange matter hypothesis, strangelets are more stable than nuclei, so nuclei are expected to decay into strangelets. But this process may be extremely slow because there is a large energy barrier to overcome: 
+                        as the weak interaction starts making a nucleus into a strangelet, the first few strange quarks form strange baryons, such as the Lambda, which are heavy. Only if many conversions occur almost simultaneously will the number of strange quarks reach the critical proportion required to achieve a lower energy state. This is very unlikely to happen, so even if the strange matter hypothesis were correct, nuclei would never be seen to decay to strangelets because their lifetime would be longer than the age of the universe.
+                        The stability of strangelets depends on their size. This is because of (a) surface tension at the interface between quark matter and vacuum (which affects small strangelets more than big ones), and (b) screening of charges, which allows small strangelets to be charged, with a neutralizing cloud of electrons/positrons around them, but requires large strangelets, like any large piece of matter, to be electrically neutral in their interior. The charge screening distance tends to be of the order of a few femtometers, so only the outer few femtometers of a strangelet can carry charge.
+                        The surface tension of strange matter is unknown. If it is smaller than a critical value (a few MeV per square femtometer) then large strangelets are unstable and will tend to fission into smaller strangelets (strange stars would still be stabilized by gravity). If it is larger than the critical value, then strangelets become more stable as they get bigger.
+                        The known particles with strange quarks are unstable. Because the strange quark is heavier than the up and down quarks, it can spontaneously decay, via the weak interaction into an up quark. Consequently particles containing strange quarks, such as the Lambda particle, always lose their strangeness, by decaying into lighter particles containing only up and down quarks.
+                        But condensed states with a larger number of quarks might not suffer from this instability. That possible stability against decay is the "strange matter hypothesis" proposed separately by Arnold Bodmer[3] and Edward Witten.[4] According to this hypothesis, when a large enough number of quarks are concentrated together, the lowest energy state is one which has roughly equal numbers of up, down, and strange quarks, namely a strangelet. This stability would occur because of the Pauli exclusion principle; having three types of quarks, rather than two as in normal nuclear matter, allows more quarks to be placed in lower energy levels
+                     """
+
+    tokenizer = AutoTokenizer.from_pretrained(torch_model.name_or_path)
+
+    xx = []
+    yy = []
+
+    for j in beam_range:
+        x = []
+        y = []
+        prev = [1, 2]
+        for i in seq_length_range:
+
+            token = tokenizer(
+                input_text,
+                padding=True,
+                truncation=True,
+                max_length=i,
+                pad_to_max_length=i,
+                return_tensors="pt",
+            )
+
+            input_ids = token["input_ids"]
+            attention_mask = token["attention_mask"]
+
+            a = pc()
+            out = onnx_model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_length=i,
+                num_beams=j,
+            )
+            b = pc()
+            x.append(b - a)
+
+            c = pc()
+            o = torch_model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_length=i,
+                num_beams=j,
+            )
+            d = pc()
+            y.append(d - c)
+
+            mean_y = np.mean(y)
+            mean_x = np.mean(x)
+            mean_ratio = mean_y / mean_x
+
+            print(f"seqL : {i}, onnx-{b-a}, pt-{d-c} .. X faster {(d-c)/(b-a)}")
+
+            # ...bleu_score-{bleu.compute(predictions=, references=[[tokenizer.decode(o.squeeze(), skip_special_tokens=True)], ])}')
+            # print(f'o---{tokenizer.decode(out.squeeze(), skip_special_tokens=True)}...p---{tokenizer.decode(o.squeeze(), skip_special_tokens=True)}')
+
+            if (o.shape[1] == prev[-1]) and (o.shape[1] == prev[-2]):
+                break
+
+            prev.append(o.shape[1])
+
+        print(f"beam no.- {j} onnx-{mean_x} pt-{mean_y} X ratio-{mean_ratio}")
+
+        xx.append(x)
+        yy.append(y)
+        plt.plot(x, "g", y, "r")
+        plt.pause(0.05)
+
+    plt.show()
+    return np.array(xx), np.array(yy)

FastT5/onnx_exporter.py

@@ -0,0 +1,294 @@
+from .huggingface_utils import get_auth_token
+from .onnx_models_structure import (
+    T5Encoder,
+    DecoderWithLMhead,
+    DecoderWithLMheadInitial,
+)
+from transformers import (
+    AutoConfig,
+    T5ForConditionalGeneration,
+    MT5ForConditionalGeneration,
+)
+import torch
+import functools
+import operator
+from progress.bar import Bar
+from pathlib import Path
+import os
+
+_folder = Path.cwd()
+saved_models_path = _folder.joinpath("models")
+
+Bar.check_tty = False
+
+
+def create_t5_encoder_decoder(pretrained_version="t5-base"):
+    """Generates an encoder and a decoder model with a language model head from a pretrained huggingface model
+
+    Args:
+        pretrained_version (str): Name of a pretrained model, or path to a pretrained / finetuned version of T5
+
+    Returns:
+        simplified_encoder: pytorch t5 encoder with a wrapper to output only the hidden states
+        decoder_with_lm_head: pytorch t5 decoder with a language modeling head
+    """
+
+    if 'mt5' in pretrained_version:
+        model = MT5ForConditionalGeneration.from_pretrained(pretrained_version, use_auth_token=get_auth_token())
+    else:
+        model = T5ForConditionalGeneration.from_pretrained(pretrained_version, use_auth_token=get_auth_token())
+
+    return turn_model_into_encoder_decoder(model)
+
+
+def turn_model_into_encoder_decoder(model):
+    encoder = model.encoder
+    decoder = model.decoder
+    lm_head = model.lm_head
+
+    decoder_with_lm_head = DecoderWithLMhead(decoder, lm_head, model.config)
+    simplified_encoder = T5Encoder(encoder)
+    decoder_with_lm_head_init = DecoderWithLMheadInitial(decoder, lm_head, model.config)
+
+    return simplified_encoder, decoder_with_lm_head, decoder_with_lm_head_init
+
+
+def generate_onnx_representation(
+    pretrained_version=None,
+    model=None,
+    output_path=None,
+    input_sequence_length=256,
+    onnx_opset_version=12,  # no other opset versions are tested, change at your own risk
+):
+    """Exports a given huggingface pretrained model, or a given model and tokenizer, to onnx
+
+    Args:
+        pretrained_version (str): Name of a pretrained model, or path to a pretrained / finetuned version of T5
+        output_path (Optional[str]): if missing then use ./models
+        input_sequence_length (Optional[int]): typical input sequence length, for use by the ORT for possible optimization
+        onnx_opset_version (Optional[int]): ONNX Operator Set Version, default 12 is the only tested version
+    """
+    if (pretrained_version is None) and model is None:
+        print(
+            "You need to specify pretrained_version (the pretrained model you wish to export). Alternatively you can export a model you have in memory."
+        )
+        return
+
+    if model is not None:
+        (
+            simplified_encoder,
+            decoder_with_lm_head,
+            decoder_with_lm_head_init,
+        ) = turn_model_into_encoder_decoder(model)
+    else:
+        (
+            simplified_encoder,
+            decoder_with_lm_head,
+            decoder_with_lm_head_init,
+        ) = create_t5_encoder_decoder(pretrained_version)
+
+    # model paths for enc, dec and dec_init
+    output_path = saved_models_path if output_path is None else Path(output_path)
+    encoder_path, decoder_path, init_decoder_path = get_model_paths(
+        pretrained_version, output_path, quantized=False
+    )
+
+    model_config = AutoConfig.from_pretrained(pretrained_version, use_auth_token=get_auth_token())
+
+    # Though these are dummy inputs, ORT optimizations do reference these values,
+    # so it is worth using values as close to production as possible
+    batch_size = 1  # not configurable since only CPU
+    enc_seq_length = input_sequence_length
+    dec_seq_length = 1  # a decoder sequence length is always one because it's just the last generated token
+    input_ids = torch.ones(batch_size, enc_seq_length, dtype=torch.int64)
+    attention_mask = torch.ones(batch_size, enc_seq_length, dtype=torch.int64)
+
+    n_heads = model_config.num_heads
+    d_kv = model_config.d_kv
+
+    input_ids_dec = torch.ones(batch_size, dec_seq_length, dtype=torch.int64)
+    attention_mask_dec = torch.ones(batch_size, dec_seq_length, dtype=torch.int64)
+    enc_out = torch.ones(
+        (batch_size, enc_seq_length, model_config.d_model), dtype=torch.float32
+    )
+
+    # self_attention_past_key_values = torch.ones(
+    #     (model_config.num_decoder_layers, 2, batch_size, n_heads, seq_length_a, d_kv), dtype=torch.float32)
+    # cross_attention_past_key_values = torch.ones(
+    #     (model_config.num_decoder_layers, 2, batch_size, n_heads, seq_length_b, d_kv), dtype=torch.float32)
+
+    sa = torch.ones(
+        (batch_size, n_heads, dec_seq_length, d_kv), dtype=torch.float32
+    )  # 1, 8, 1, 64
+    ca = torch.ones(
+        (batch_size, n_heads, enc_seq_length, d_kv), dtype=torch.float32
+    )  # 1, 8, variable, 64
+    t5_block = (sa, sa, ca, ca)
+    past_key_values = (t5_block,) * model_config.num_decoder_layers
+
+    flat_past_key_values = functools.reduce(operator.iconcat, past_key_values, [])
+
+    decoder_all_inputs = tuple(
+        [input_ids_dec, attention_mask_dec, enc_out] + flat_past_key_values
+    )
+
+    # for progress bars
+    bar = Bar("Exporting to onnx...", max=3)
+
+    import warnings
+
+    # ignores all the warnings during conversion
+    warnings.filterwarnings("ignore")
+
+    # Exports to ONNX
+    with torch.no_grad():
+
+        decoder_inputs = [
+            "input_ids",
+            "encoder_attention_mask",
+            "encoder_hidden_states",
+        ]
+
+        pkv_input_names = ["pkv_{}".format(i) for i in range(len(flat_past_key_values))]
+
+        decoder_input_names = decoder_inputs + pkv_input_names
+
+        decoder_output_names = ["logits", "output_past_key_values"]
+
+        dyn_axis_general = {0: "batch", 1: "sequence"}
+        dyn_axis_pkv = {0: "batch", 2: "seq_length"}
+
+        dyn_axis = {
+            "input_ids": dyn_axis_general,
+            "encoder_attention_mask": dyn_axis_general,
+            "encoder_hidden_states": dyn_axis_general,
+            "logits": dyn_axis_general,
+            "output_past_key_values": dyn_axis_general,
+        }
+
+        dyn_pkv = {
+            "pkv_{}".format(i): dyn_axis_pkv
+            for i in range(len(flat_past_key_values))
+        }
+
+        dyn_axis_params = {**dyn_axis, **dyn_pkv}
+
+        # decoder to utilize past key values:
+        torch.onnx.export(
+            decoder_with_lm_head,
+            decoder_all_inputs,
+            decoder_path.as_posix(),
+            export_params=True,
+            do_constant_folding=True,
+            opset_version=onnx_opset_version,
+            input_names=decoder_input_names,
+            output_names=decoder_output_names,
+            dynamic_axes=dyn_axis_params,
+        )
+        bar.next()
+
+        torch.onnx.export(
+            simplified_encoder,
+            args=(input_ids, attention_mask),
+            f=encoder_path.as_posix(),
+            export_params=True,
+            opset_version=onnx_opset_version,
+            do_constant_folding=True,
+            input_names=["input_ids", "attention_mask"],
+            output_names=["hidden_states"],
+            dynamic_axes={
+                "input_ids": dyn_axis_general,
+                "attention_mask": dyn_axis_general,
+                "hidden_states": dyn_axis_general,
+            },
+        )
+        bar.next()
+        # initial decoder to produce past key values
+        torch.onnx.export(
+            decoder_with_lm_head_init,
+            (input_ids_dec, attention_mask_dec, enc_out),
+            init_decoder_path.as_posix(),
+            export_params=True,
+            opset_version=onnx_opset_version,
+            input_names=[
+                "input_ids",
+                "encoder_attention_mask",
+                "encoder_hidden_states",
+            ],
+            output_names=["logits", "past_key_values"],
+            dynamic_axes={
+                # batch_size, seq_length = input_shape
+                "input_ids": dyn_axis_general,
+                "encoder_attention_mask": dyn_axis_general,
+                "encoder_hidden_states": dyn_axis_general,
+                "logits": dyn_axis_general,
+                "past_key_values": dyn_axis_general,
+            },
+        )
+        bar.next()
+        bar.finish()
+
+    return encoder_path, decoder_path, init_decoder_path
+
+
+def get_model_paths(pretrained_model, model_path, quantized):
+
+    model_path.mkdir(parents=True, exist_ok=True)
+
+    # gets only the filename
+    pretrained_model_name = Path(pretrained_model).stem
+
+    if not quantized:
+        encoder_path = model_path.joinpath(f"{pretrained_model_name}-encoder.onnx")
+        decoder_path = model_path.joinpath(f"{pretrained_model_name}-decoder.onnx")
+        init_decoder_path = model_path.joinpath(
+            f"{pretrained_model_name}-init-decoder.onnx"
+        )
+    else:
+        encoder_path = model_path.joinpath(
+            f"{pretrained_model_name}-encoder-quantized.onnx"
+        )
+        decoder_path = model_path.joinpath(
+            f"{pretrained_model_name}-decoder-quantized.onnx"
+        )
+        init_decoder_path = model_path.joinpath(
+            f"{pretrained_model_name}-init-decoder-quantized.onnx"
+        )
+
+    return encoder_path, decoder_path, init_decoder_path
+
+
+def quantize(models_name_or_path):
+    """
+    Quantize the weights of the model from float32 to in8 to allow very efficient inference on modern CPU
+
+    Uses unsigned ints for activation values, signed ints for weights, per
+    https://onnxruntime.ai/docs/performance/quantization.html#data-type-selection
+    it is faster on most CPU architectures
+    Args:
+        onnx_model_path: Path to location the exported ONNX model is stored
+    Returns: The Path generated for the quantized
+    """
+    from onnxruntime.quantization import quantize_dynamic, QuantType
+
+    bar = Bar("Quantizing...", max=3)
+
+    quant_model_paths = []
+    for model in models_name_or_path:
+        model_name = model.as_posix()
+        output_model_name = f"{model_name[:-5]}-quantized.onnx"
+        quantize_dynamic(
+            model_input=model_name,
+            model_output=output_model_name,
+            per_channel=True,
+            reduce_range=True, # should be the same as per_channel
+            activation_type=QuantType.QUInt8,
+            weight_type=QuantType.QInt8,  # per docs, signed is faster on most CPUs
+            optimize_model=False,
+        )  # op_types_to_quantize=['MatMul', 'Relu', 'Add', 'Mul' ],
+        quant_model_paths.append(output_model_name)
+        bar.next()
+
+    bar.finish()
+
+    return tuple(quant_model_paths)

FastT5/onnx_models.py

@@ -0,0 +1,269 @@
+from .huggingface_utils import get_auth_token
+from .ort_settings import get_onnx_runtime_sessions
+from .onnx_exporter import (
+    generate_onnx_representation,
+    quantize,
+    get_model_paths,
+    saved_models_path,
+)
+from pathlib import Path
+
+from transformers import (
+    AutoConfig,
+    MT5Config,
+    T5ForConditionalGeneration,
+)
+from transformers.modeling_outputs import (
+    Seq2SeqLMOutput,
+    BaseModelOutput,
+)
+import torch
+import functools
+import operator
+import numpy
+
+
+class T5Encoder(torch.nn.Module):
+    def __init__(self, encoder_sess):
+        super().__init__()
+        self.encoder = encoder_sess
+        self.main_input_name = "input_ids"
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask,
+        inputs_embeds=None,
+        head_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+
+        encoder_hidden_state = torch.from_numpy(
+            self.encoder.run(
+                None,
+                {
+                    "input_ids": input_ids.cpu().numpy(),
+                    "attention_mask": attention_mask.cpu().numpy(),
+                },
+            )[0]
+        )
+
+        return BaseModelOutput(encoder_hidden_state)
+
+
+class T5DecoderInit(torch.nn.Module):
+    def __init__(self, decoder_sess):
+        super().__init__()
+        self.decoder = decoder_sess
+
+    def forward(self, input_ids, encoder_attention_mask, encoder_hidden_states):
+
+        decoder_outputs = self.decoder.run(
+            None,
+            {
+                "input_ids": input_ids.cpu().numpy(),
+                "encoder_attention_mask": encoder_attention_mask.cpu().numpy(),
+                "encoder_hidden_states": encoder_hidden_states.cpu().numpy(),
+            },
+        )
+
+        list_pkv = tuple(torch.from_numpy(x) for x in decoder_outputs[1:])
+
+        out_past_key_values = tuple(
+            list_pkv[i: i + 4] for i in range(0, len(list_pkv), 4)
+        )
+
+        return torch.from_numpy(decoder_outputs[0]), out_past_key_values
+
+
+class T5Decoder(torch.nn.Module):
+    def __init__(self, decoder_sess):
+        super().__init__()
+        self.decoder = decoder_sess
+
+    def forward(self, input_ids, attention_mask, encoder_output, past_key_values):
+
+        decoder_inputs = {
+            "input_ids": input_ids.cpu().numpy(),
+            "encoder_attention_mask": attention_mask.cpu().numpy(),
+            "encoder_hidden_states": encoder_output.cpu().numpy(),
+        }
+
+        flat_past_key_values = functools.reduce(
+            operator.iconcat, past_key_values, [])
+
+        past_key_values = {
+            f"pkv_{i}": pkv.cpu().numpy() for i, pkv in enumerate(flat_past_key_values)
+        }
+
+        decoder_outputs = self.decoder.run(
+            None, {**decoder_inputs, **past_key_values})
+        # converts each value of the list to tensor from numpy
+        list_pkv = tuple(torch.from_numpy(x) for x in decoder_outputs[1:])
+
+        # creates a tuple of tuples of shape 6x4 from the above tuple
+        out_past_key_values = tuple(
+            list_pkv[i: i + 4] for i in range(0, len(list_pkv), 4)
+        )
+
+        return torch.from_numpy(decoder_outputs[0]), out_past_key_values
+
+
+class OnnxT5(T5ForConditionalGeneration):
+    """creates a T5 model using onnx sessions (encode, decoder & init_decoder)"""
+
+    def __init__(self, model_or_model_path, onnx_model_sessions):
+        config = AutoConfig.from_pretrained(
+            model_or_model_path, use_auth_token=get_auth_token()
+        )
+        super().__init__(config)
+
+        # monkeypatch to work for MT5
+        if (
+            isinstance(model_or_model_path, str)
+            and "mt5" in model_or_model_path.lower()
+        ) or (
+            hasattr(model_or_model_path, "name_or_path")
+            and "mt5" in model_or_model_path.name_or_path
+        ):
+            self.model_type = "mt5"
+            self.config_class = MT5Config
+            self._keys_to_ignore_on_load_missing = [
+                r"encoder\.embed_tokens\.weight",
+            ]
+            self._keys_to_ignore_on_save = [
+                r"encoder\.embed_tokens\.weight",
+            ]
+
+        assert len(onnx_model_sessions) == 3, "all three models should be given"
+
+        encoder_sess, decoder_sess, decoder_sess_init = onnx_model_sessions
+
+        self.encoder = T5Encoder(encoder_sess)
+        self.decoder = T5Decoder(decoder_sess)
+        self.decoder_init = T5DecoderInit(decoder_sess_init)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        decoder_input_ids=None,
+        decoder_attention_mask=None,
+        head_mask=None,
+        decoder_head_mask=None,
+        cross_attn_head_mask=None,
+        encoder_outputs=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        decoder_inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+
+        if encoder_outputs is None:
+            # Convert encoder inputs in embeddings if needed
+            encoder_outputs = self.encoder(
+                input_ids=input_ids, attention_mask=attention_mask
+            )
+
+        encoder_hidden_states = encoder_outputs[0]
+
+        if past_key_values is not None:
+            if decoder_input_ids is not None:
+                decoder_input_ids = decoder_input_ids[:, -1:]
+            if decoder_inputs_embeds is not None:
+                decoder_inputs_embeds = decoder_inputs_embeds[:, -1:]
+
+        if past_key_values is None:
+
+            # runs only for the first time:
+            init_onnx_outputs = self.decoder_init(
+                decoder_input_ids, attention_mask, encoder_hidden_states
+            )
+
+            logits, past_key_values = init_onnx_outputs
+
+        else:
+
+            onnx_outputs = self.decoder(
+                decoder_input_ids,
+                attention_mask,
+                encoder_hidden_states,
+                past_key_values,
+            )
+
+            logits, past_key_values = onnx_outputs
+
+        return Seq2SeqLMOutput(logits=logits, past_key_values=past_key_values)
+
+
+def export_and_get_onnx_model(
+    model_or_model_path, custom_output_path=saved_models_path, quantized=True
+):
+    """
+                          Method for whole pipeline,
+    converts from pytorch to onnx --> quantizes model --> sets onnx runtime
+                --> builds whole onnx model with all sessions
+
+    """
+
+    # Step 1. convert huggingfaces t5 model to onnx
+    onnx_model_paths = generate_onnx_representation(
+        model_or_model_path, output_path=custom_output_path
+    )
+
+    if quantized:
+        # Step 2. (recommended) quantize the converted model for fast inference and to reduce model size.
+        quant_model_paths = quantize(onnx_model_paths)
+
+        # step 3. setup onnx runtime
+        print("Setting up onnx model...")
+        model_sessions = get_onnx_runtime_sessions(quant_model_paths)
+    else:
+        print("Setting up onnx model...")
+        model_sessions = get_onnx_runtime_sessions(onnx_model_paths)
+
+    # step 4. get the onnx model
+    model = OnnxT5(model_or_model_path, model_sessions)
+    print("Done!")
+
+    return model
+
+
+def get_onnx_model(model_name, onnx_models_path=saved_models_path, quantized=True):
+    """
+    method gets the onnx model, if already converted models exists
+    Example:
+    >> get_onnx_model(model_name="t5-finetuned", onnx_models_path="../models/onnx/quantized/")
+
+    """
+
+    encoder_path, decoder_path, init_decoder_path = get_model_paths(
+        model_name, Path(onnx_models_path), quantized
+    )
+
+    if quantized:
+        assert (
+            encoder_path.exists()
+            and decoder_path.exists()
+            and init_decoder_path.exists()
+        ), "quantized model don't exist in the model folder, first quantize the model!"
+    else:
+        assert (
+            encoder_path.exists()
+            and decoder_path.exists()
+            and init_decoder_path.exists()
+        ), "all or some models don't exists in the model folder, first convert the model! "
+
+    model_paths = encoder_path, decoder_path, init_decoder_path
+
+    model_sessions = get_onnx_runtime_sessions(model_paths)
+
+    model = OnnxT5(model_name, model_sessions)
+
+    return model

FastT5/onnx_models_structure.py

@@ -0,0 +1,62 @@
+import torch
+
+
+class DecoderWithLMhead(torch.nn.Module):
+    """ Creation of a class to combine the decoder and the lm head """
+
+    def __init__(self, decoder, lm_head, config):
+        super().__init__()
+        self.decoder = decoder
+        self.lm_head = lm_head
+        self.config = config
+
+    def forward(self, *inputs):
+
+        input_ids, attention_mask, encoder_hidden_states = inputs[:3]
+
+        list_pkv = inputs[3:]
+        past_key_values = tuple(list_pkv[i : i + 4] for i in range(0, len(list_pkv), 4))
+
+        decoder_output = self.decoder(
+            input_ids=input_ids,  # decoder_input_ids
+            encoder_attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            past_key_values=past_key_values,
+        )
+
+        lm_head_out = self.lm_head(decoder_output[0] * (self.config.d_model ** -0.5))
+
+        return lm_head_out, decoder_output[1]
+
+
+class T5Encoder(torch.nn.Module):
+    """ Creation of a class to output only the last hidden state from the encoder """
+
+    def __init__(self, encoder):
+        super().__init__()
+        self.encoder = encoder
+
+    def forward(self, *input, **kwargs):
+        return self.encoder(*input, **kwargs)[0]
+
+
+class DecoderWithLMheadInitial(torch.nn.Module):
+    """ Creation of a class to combine the decoder and the lm head """
+
+    def __init__(self, decoder, lm_head, config):
+        super().__init__()
+        self.decoder = decoder
+        self.lm_head = lm_head
+        self.config = config
+
+    def forward(self, input_ids, attention_mask, encoder_hidden_states):
+        decoder_output = self.decoder(
+            input_ids=input_ids,
+            encoder_attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+        )
+
+        return (
+            self.lm_head(decoder_output[0] * (self.config.d_model ** -0.5)),
+            decoder_output[1],
+        )

FastT5/ort_settings.py

@@ -0,0 +1,96 @@
+import os, psutil
+
+os.environ["OMP_NUM_THREADS"] = str(psutil.cpu_count(logical=True))
+os.environ["OMP_WAIT_POLICY"] = "ACTIVE"
+
+
+from onnxruntime import (
+    GraphOptimizationLevel,
+    InferenceSession,
+    SessionOptions,
+    ExecutionMode,
+)
+
+
+def get_onnx_runtime_sessions(
+    model_paths,
+    default: bool = True,
+    opt_level: int = 99,
+    parallel_exe_mode: bool = True,
+    n_threads: int = 0,
+    provider=[
+        "CPUExecutionProvider",
+    ],
+) -> InferenceSession:
+    """
+            Optimizes the model
+
+    Args:
+        model_paths (List or Tuple of str) : the path to, in order:
+            path_to_encoder (str) : the path of input onnx encoder model.
+            path_to_decoder (str) : the path of input onnx decoder model.
+            path_to_initial_decoder (str) :  the path of input initial onnx decoder model.
+        default : set this to true, ort will choose the best settings for your hardware.
+                  (you can test out different settings for better results.)
+        opt_level (int) : sess_options.GraphOptimizationLevel param if set 1 uses 'ORT_ENABLE_BASIC',
+                          2 for 'ORT_ENABLE_EXTENDED' and 99 for 'ORT_ENABLE_ALL',
+                          default value is set to 99.
+        parallel_exe_mode (bool) :  Sets the execution mode. Default is True (parallel).
+        n_threads (int) :  Sets the number of threads used to parallelize the execution within nodes. Default is 0 to let onnxruntime choose
+        provider : execution providers list.
+
+    Returns:
+        encoder_session : encoder onnx InferenceSession
+        decoder_session : decoder onnx InferenceSession
+        decoder_sess_init : initial decoder onnx InferenceSession
+
+    """
+    path_to_encoder, path_to_decoder, path_to_initial_decoder = model_paths
+
+    if default:
+
+        encoder_sess = InferenceSession(str(path_to_encoder))
+
+        decoder_sess = InferenceSession(str(path_to_decoder))
+
+        decoder_sess_init = InferenceSession(str(path_to_initial_decoder))
+
+    else:
+
+        # Few properties that might have an impact on performances
+        options = SessionOptions()
+
+        if opt_level == 1:
+            options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC
+        elif opt_level == 2:
+            options.graph_optimization_level = (
+                GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+            )
+        else:
+            assert opt_level == 99
+            options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        # set this true for better performance
+        if parallel_exe_mode == True:
+            options.execution_mode = ExecutionMode.ORT_PARALLEL
+        else:
+            options.execution_mode = ExecutionMode.ORT_SEQUENTIAL
+
+        options.intra_op_num_threads = n_threads
+        # options.inter_op_num_threads = 10
+
+        # options.enable_profiling = True
+
+        encoder_sess = InferenceSession(
+            str(path_to_encoder), options, providers=provider
+        )
+
+        decoder_sess = InferenceSession(
+            str(path_to_decoder), options, providers=provider
+        )
+
+        decoder_sess_init = InferenceSession(
+            str(path_to_initial_decoder), options, providers=provider
+        )
+
+    return encoder_sess, decoder_sess, decoder_sess_init

app.py

@@ -1,742 +1,9 @@
-import psutil
-from transformers import (
-    AutoConfig,
-    T5ForConditionalGeneration,
-    MT5ForConditionalGeneration,
-)
-import torch
 import time
 import gradio as gr
 from transformers import AutoTokenizer
-import onnxruntime as ort
-from transformers.modeling_outputs import (
-    Seq2SeqLMOutput,
-    BaseModelOutput,
-)
 import os
 from pathlib import Path
-from progress.bar import Bar
-import operator
-import functools
-from onnxruntime import (
-    GraphOptimizationLevel,
-    InferenceSession,
-    SessionOptions,
-    ExecutionMode,
-)
-_auth_token = None
-
-
-def set_auth_token(token):
-    """Set the token which allows the user to authenticate to hugginface.co for downloading private models
-
-    Args:
-        token (Union[str, bool]): The token value to store. One of:
-            - an API key (from https://huggingface.co/organizations/ORGNAME/settings/token),
-            - a login token obtained by running `$ transformers-cli login`
-            - `True`, which tells transformers to use the login token stored in ~/.huggingface/token
-
-    Returns:
-        None
-    """
-    global _auth_token
-    _auth_token = token
-
-
-def get_auth_token():
-    """Get the user-configurable auth token, which defaults to None
-
-    Returns:
-        auth_token (Optional[Union[str, bool]]) for authenticating with huggingface.co
-    """
-    global _auth_token
-    return _auth_token
-
-
-os.environ["OMP_NUM_THREADS"] = str(psutil.cpu_count(logical=True))
-os.environ["OMP_WAIT_POLICY"] = "ACTIVE"
-
-
-def get_onnx_runtime_sessions(
-    model_paths,
-    default: bool = True,
-    opt_level: int = 99,
-    parallel_exe_mode: bool = True,
-    n_threads: int = 0,
-    provider=[
-        "CPUExecutionProvider",
-    ],
-) -> InferenceSession:
-    """
-            Optimizes the model
-
-    Args:
-        model_paths (List or Tuple of str) : the path to, in order:
-            path_to_encoder (str) : the path of input onnx encoder model.
-            path_to_decoder (str) : the path of input onnx decoder model.
-            path_to_initial_decoder (str) :  the path of input initial onnx decoder model.
-        default : set this to true, ort will choose the best settings for your hardware.
-                  (you can test out different settings for better results.)
-        opt_level (int) : sess_options.GraphOptimizationLevel param if set 1 uses 'ORT_ENABLE_BASIC',
-                          2 for 'ORT_ENABLE_EXTENDED' and 99 for 'ORT_ENABLE_ALL',
-                          default value is set to 99.
-        parallel_exe_mode (bool) :  Sets the execution mode. Default is True (parallel).
-        n_threads (int) :  Sets the number of threads used to parallelize the execution within nodes. Default is 0 to let onnxruntime choose
-        provider : execution providers list.
-
-    Returns:
-        encoder_session : encoder onnx InferenceSession
-        decoder_session : decoder onnx InferenceSession
-        decoder_sess_init : initial decoder onnx InferenceSession
-
-    """
-    path_to_encoder, path_to_decoder, path_to_initial_decoder = model_paths
-
-    if default:
-
-        encoder_sess = InferenceSession(str(path_to_encoder))
-
-        decoder_sess = InferenceSession(str(path_to_decoder))
-
-        decoder_sess_init = InferenceSession(str(path_to_initial_decoder))
-
-    else:
-
-        # Few properties that might have an impact on performances
-        options = SessionOptions()
-
-        if opt_level == 1:
-            options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC
-        elif opt_level == 2:
-            options.graph_optimization_level = (
-                GraphOptimizationLevel.ORT_ENABLE_EXTENDED
-            )
-        else:
-            assert opt_level == 99
-            options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
-
-        # set this true for better performance
-        if parallel_exe_mode == True:
-            options.execution_mode = ExecutionMode.ORT_PARALLEL
-        else:
-            options.execution_mode = ExecutionMode.ORT_SEQUENTIAL
-
-        options.intra_op_num_threads = n_threads
-        # options.inter_op_num_threads = 10
-
-        # options.enable_profiling = True
-
-        encoder_sess = InferenceSession(
-            str(path_to_encoder), options, providers=provider
-        )
-
-        decoder_sess = InferenceSession(
-            str(path_to_decoder), options, providers=provider
-        )
-
-        decoder_sess_init = InferenceSession(
-            str(path_to_initial_decoder), options, providers=provider
-        )
-
-    return encoder_sess, decoder_sess, decoder_sess_init
-
-
-class DecoderWithLMhead(torch.nn.Module):
-    """ Creation of a class to combine the decoder and the lm head """
-
-    def __init__(self, decoder, lm_head, config):
-        super().__init__()
-        self.decoder = decoder
-        self.lm_head = lm_head
-        self.config = config
-
-    def forward(self, *inputs):
-
-        input_ids, attention_mask, encoder_hidden_states = inputs[:3]
-
-        list_pkv = inputs[3:]
-        past_key_values = tuple(list_pkv[i: i + 4]
-                                for i in range(0, len(list_pkv), 4))
-
-        decoder_output = self.decoder(
-            input_ids=input_ids,  # decoder_input_ids
-            encoder_attention_mask=attention_mask,
-            encoder_hidden_states=encoder_hidden_states,
-            past_key_values=past_key_values,
-        )
-
-        lm_head_out = self.lm_head(
-            decoder_output[0] * (self.config.d_model ** -0.5))
-
-        return lm_head_out, decoder_output[1]
-
-
-class T5Encoder(torch.nn.Module):
-    """ Creation of a class to output only the last hidden state from the encoder """
-
-    def __init__(self, encoder):
-        super().__init__()
-        self.encoder = encoder
-
-    def forward(self, *input, **kwargs):
-        return self.encoder(*input, **kwargs)[0]
-
-
-class DecoderWithLMheadInitial(torch.nn.Module):
-    """ Creation of a class to combine the decoder and the lm head """
-
-    def __init__(self, decoder, lm_head, config):
-        super().__init__()
-        self.decoder = decoder
-        self.lm_head = lm_head
-        self.config = config
-
-    def forward(self, input_ids, attention_mask, encoder_hidden_states):
-        decoder_output = self.decoder(
-            input_ids=input_ids,
-            encoder_attention_mask=attention_mask,
-            encoder_hidden_states=encoder_hidden_states,
-        )
-
-        return (
-            self.lm_head(decoder_output[0] * (self.config.d_model ** -0.5)),
-            decoder_output[1],
-        )
-
-
-_folder = Path.cwd()
-saved_models_path = _folder.joinpath("models")
-
-Bar.check_tty = False
-
-
-def create_t5_encoder_decoder(pretrained_version="t5-base"):
-    """Generates an encoder and a decoder model with a language model head from a pretrained huggingface model
-
-    Args:
-        pretrained_version (str): Name of a pretrained model, or path to a pretrained / finetuned version of T5
-
-    Returns:
-        simplified_encoder: pytorch t5 encoder with a wrapper to output only the hidden states
-        decoder_with_lm_head: pytorch t5 decoder with a language modeling head
-    """
-
-    if 'mt5' in pretrained_version:
-        model = MT5ForConditionalGeneration.from_pretrained(
-            pretrained_version, use_auth_token=get_auth_token())
-    else:
-        model = T5ForConditionalGeneration.from_pretrained(
-            pretrained_version, use_auth_token=get_auth_token())
-
-    return turn_model_into_encoder_decoder(model)
-
-
-def turn_model_into_encoder_decoder(model):
-    encoder = model.encoder
-    decoder = model.decoder
-    lm_head = model.lm_head
-
-    decoder_with_lm_head = DecoderWithLMhead(decoder, lm_head, model.config)
-    simplified_encoder = T5Encoder(encoder)
-    decoder_with_lm_head_init = DecoderWithLMheadInitial(
-        decoder, lm_head, model.config)
-
-    return simplified_encoder, decoder_with_lm_head, decoder_with_lm_head_init
-
-
-def generate_onnx_representation(
-    pretrained_version=None,
-    model=None,
-    output_path=None,
-    input_sequence_length=256,
-    onnx_opset_version=12,  # no other opset versions are tested, change at your own risk
-):
-    """Exports a given huggingface pretrained model, or a given model and tokenizer, to onnx
-
-    Args:
-        pretrained_version (str): Name of a pretrained model, or path to a pretrained / finetuned version of T5
-        output_path (Optional[str]): if missing then use ./models
-        input_sequence_length (Optional[int]): typical input sequence length, for use by the ORT for possible optimization
-        onnx_opset_version (Optional[int]): ONNX Operator Set Version, default 12 is the only tested version
-    """
-    if (pretrained_version is None) and model is None:
-        print(
-            "You need to specify pretrained_version (the pretrained model you wish to export). Alternatively you can export a model you have in memory."
-        )
-        return
-
-    if model is not None:
-        (
-            simplified_encoder,
-            decoder_with_lm_head,
-            decoder_with_lm_head_init,
-        ) = turn_model_into_encoder_decoder(model)
-    else:
-        (
-            simplified_encoder,
-            decoder_with_lm_head,
-            decoder_with_lm_head_init,
-        ) = create_t5_encoder_decoder(pretrained_version)
-
-    # model paths for enc, dec and dec_init
-    output_path = saved_models_path if output_path is None else Path(
-        output_path)
-    encoder_path, decoder_path, init_decoder_path = get_model_paths(
-        pretrained_version, output_path, quantized=False
-    )
-
-    model_config = AutoConfig.from_pretrained(
-        pretrained_version, use_auth_token=get_auth_token())
-
-    # Though these are dummy inputs, ORT optimizations do reference these values,
-    # so it is worth using values as close to production as possible
-    batch_size = 1  # not configurable since only CPU
-    enc_seq_length = input_sequence_length
-    # a decoder sequence length is always one because it's just the last generated token
-    dec_seq_length = 1
-    input_ids = torch.ones(batch_size, enc_seq_length, dtype=torch.int64)
-    attention_mask = torch.ones(batch_size, enc_seq_length, dtype=torch.int64)
-
-    n_heads = model_config.num_heads
-    d_kv = model_config.d_kv
-
-    input_ids_dec = torch.ones(batch_size, dec_seq_length, dtype=torch.int64)
-    attention_mask_dec = torch.ones(
-        batch_size, dec_seq_length, dtype=torch.int64)
-    enc_out = torch.ones(
-        (batch_size, enc_seq_length, model_config.d_model), dtype=torch.float32
-    )
-
-    # self_attention_past_key_values = torch.ones(
-    #     (model_config.num_decoder_layers, 2, batch_size, n_heads, seq_length_a, d_kv), dtype=torch.float32)
-    # cross_attention_past_key_values = torch.ones(
-    #     (model_config.num_decoder_layers, 2, batch_size, n_heads, seq_length_b, d_kv), dtype=torch.float32)
-
-    sa = torch.ones(
-        (batch_size, n_heads, dec_seq_length, d_kv), dtype=torch.float32
-    )  # 1, 8, 1, 64
-    ca = torch.ones(
-        (batch_size, n_heads, enc_seq_length, d_kv), dtype=torch.float32
-    )  # 1, 8, variable, 64
-    t5_block = (sa, sa, ca, ca)
-    past_key_values = (t5_block,) * model_config.num_decoder_layers
-
-    flat_past_key_values = functools.reduce(
-        operator.iconcat, past_key_values, [])
-
-    decoder_all_inputs = tuple(
-        [input_ids_dec, attention_mask_dec, enc_out] + flat_past_key_values
-    )
-
-    # for progress bars
-    bar = Bar("Exporting to onnx...", max=3)
-
-    import warnings
-
-    # ignores all the warnings during conversion
-    warnings.filterwarnings("ignore")
-
-    # Exports to ONNX
-    with torch.no_grad():
-
-        decoder_inputs = [
-            "input_ids",
-            "encoder_attention_mask",
-            "encoder_hidden_states",
-        ]
-
-        pkv_input_names = ["pkv_{}".format(
-            i) for i in range(len(flat_past_key_values))]
-
-        decoder_input_names = decoder_inputs + pkv_input_names
-
-        decoder_output_names = ["logits", "output_past_key_values"]
-
-        dyn_axis_general = {0: "batch", 1: "sequence"}
-        dyn_axis_pkv = {0: "batch", 2: "seq_length"}
-
-        dyn_axis = {
-            "input_ids": dyn_axis_general,
-            "encoder_attention_mask": dyn_axis_general,
-            "encoder_hidden_states": dyn_axis_general,
-            "logits": dyn_axis_general,
-            "output_past_key_values": dyn_axis_general,
-        }
-
-        dyn_pkv = {
-            "pkv_{}".format(i): dyn_axis_pkv
-            for i in range(len(flat_past_key_values))
-        }
-
-        dyn_axis_params = {**dyn_axis, **dyn_pkv}
-
-        # decoder to utilize past key values:
-        torch.onnx.export(
-            decoder_with_lm_head,
-            decoder_all_inputs,
-            decoder_path.as_posix(),
-            export_params=True,
-            do_constant_folding=True,
-            opset_version=onnx_opset_version,
-            input_names=decoder_input_names,
-            output_names=decoder_output_names,
-            dynamic_axes=dyn_axis_params,
-        )
-        bar.next()
-
-        torch.onnx.export(
-            simplified_encoder,
-            args=(input_ids, attention_mask),
-            f=encoder_path.as_posix(),
-            export_params=True,
-            opset_version=onnx_opset_version,
-            do_constant_folding=True,
-            input_names=["input_ids", "attention_mask"],
-            output_names=["hidden_states"],
-            dynamic_axes={
-                "input_ids": dyn_axis_general,
-                "attention_mask": dyn_axis_general,
-                "hidden_states": dyn_axis_general,
-            },
-        )
-        bar.next()
-        # initial decoder to produce past key values
-        torch.onnx.export(
-            decoder_with_lm_head_init,
-            (input_ids_dec, attention_mask_dec, enc_out),
-            init_decoder_path.as_posix(),
-            export_params=True,
-            opset_version=onnx_opset_version,
-            input_names=[
-                "input_ids",
-                "encoder_attention_mask",
-                "encoder_hidden_states",
-            ],
-            output_names=["logits", "past_key_values"],
-            dynamic_axes={
-                # batch_size, seq_length = input_shape
-                "input_ids": dyn_axis_general,
-                "encoder_attention_mask": dyn_axis_general,
-                "encoder_hidden_states": dyn_axis_general,
-                "logits": dyn_axis_general,
-                "past_key_values": dyn_axis_general,
-            },
-        )
-        bar.next()
-        bar.finish()
-
-    return encoder_path, decoder_path, init_decoder_path
-
-
-def get_model_paths(pretrained_model, model_path, quantized):
-
-    model_path.mkdir(parents=True, exist_ok=True)
-
-    # gets only the filename
-    pretrained_model_name = Path(pretrained_model).stem
-
-    if not quantized:
-        encoder_path = model_path.joinpath(
-            f"{pretrained_model_name}-encoder.onnx")
-        decoder_path = model_path.joinpath(
-            f"{pretrained_model_name}-decoder.onnx")
-        init_decoder_path = model_path.joinpath(
-            f"{pretrained_model_name}-init-decoder.onnx"
-        )
-    else:
-        encoder_path = model_path.joinpath(
-            f"{pretrained_model_name}-encoder-quantized.onnx"
-        )
-        decoder_path = model_path.joinpath(
-            f"{pretrained_model_name}-decoder-quantized.onnx"
-        )
-        init_decoder_path = model_path.joinpath(
-            f"{pretrained_model_name}-init-decoder-quantized.onnx"
-        )
-
-    return encoder_path, decoder_path, init_decoder_path
-
-
-def quantize(models_name_or_path):
-    """
-    Quantize the weights of the model from float32 to in8 to allow very efficient inference on modern CPU
-
-    Uses unsigned ints for activation values, signed ints for weights, per
-    https://onnxruntime.ai/docs/performance/quantization.html#data-type-selection
-    it is faster on most CPU architectures
-    Args:
-        onnx_model_path: Path to location the exported ONNX model is stored
-    Returns: The Path generated for the quantized
-    """
-    from onnxruntime.quantization import quantize_dynamic, QuantType
-
-    bar = Bar("Quantizing...", max=3)
-
-    quant_model_paths = []
-    for model in models_name_or_path:
-        model_name = model.as_posix()
-        output_model_name = f"{model_name[:-5]}-quantized.onnx"
-        quantize_dynamic(
-            model_input=model_name,
-            model_output=output_model_name,
-            per_channel=True,
-            reduce_range=True,  # should be the same as per_channel
-            activation_type=QuantType.QUInt8,
-            weight_type=QuantType.QInt8,  # per docs, signed is faster on most CPUs
-            optimize_model=False,
-        )  # op_types_to_quantize=['MatMul', 'Relu', 'Add', 'Mul' ],
-        quant_model_paths.append(output_model_name)
-        bar.next()
-
-    bar.finish()
-
-    return tuple(quant_model_paths)
-
-
-class T5Encoder(torch.nn.Module):
-    def __init__(self, encoder_sess):
-        super().__init__()
-        self.encoder = encoder_sess
-        self.main_input_name = "input_ids"
-
-    def forward(
-        self,
-        input_ids,
-        attention_mask,
-        inputs_embeds=None,
-        head_mask=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        return_dict=None,
-    ):
-
-        encoder_hidden_state = torch.from_numpy(
-            self.encoder.run(
-                None,
-                {
-                    "input_ids": input_ids.cpu().numpy(),
-                    "attention_mask": attention_mask.cpu().numpy(),
-                },
-            )[0]
-        )
-
-        return BaseModelOutput(encoder_hidden_state)
-
-
-class T5DecoderInit(torch.nn.Module):
-    def __init__(self, decoder_sess):
-        super().__init__()
-        self.decoder = decoder_sess
-
-    def forward(self, input_ids, encoder_attention_mask, encoder_hidden_states):
-
-        decoder_outputs = self.decoder.run(
-            None,
-            {
-                "input_ids": input_ids.cpu().numpy(),
-                "encoder_attention_mask": encoder_attention_mask.cpu().numpy(),
-                "encoder_hidden_states": encoder_hidden_states.cpu().numpy(),
-            },
-        )
-
-        list_pkv = tuple(torch.from_numpy(x) for x in decoder_outputs[1:])
-
-        out_past_key_values = tuple(
-            list_pkv[i: i + 4] for i in range(0, len(list_pkv), 4)
-        )
-
-        return torch.from_numpy(decoder_outputs[0]), out_past_key_values
-
-
-class T5Decoder(torch.nn.Module):
-    def __init__(self, decoder_sess):
-        super().__init__()
-        self.decoder = decoder_sess
-
-    def forward(self, input_ids, attention_mask, encoder_output, past_key_values):
-
-        decoder_inputs = {
-            "input_ids": input_ids.cpu().numpy(),
-            "encoder_attention_mask": attention_mask.cpu().numpy(),
-            "encoder_hidden_states": encoder_output.cpu().numpy(),
-        }
-
-        flat_past_key_values = functools.reduce(
-            operator.iconcat, past_key_values, [])
-
-        past_key_values = {
-            f"pkv_{i}": pkv.cpu().numpy() for i, pkv in enumerate(flat_past_key_values)
-        }
-
-        decoder_outputs = self.decoder.run(
-            None, {**decoder_inputs, **past_key_values})
-        # converts each value of the list to tensor from numpy
-        list_pkv = tuple(torch.from_numpy(x) for x in decoder_outputs[1:])
-
-        # creates a tuple of tuples of shape 6x4 from the above tuple
-        out_past_key_values = tuple(
-            list_pkv[i: i + 4] for i in range(0, len(list_pkv), 4)
-        )
-
-        return torch.from_numpy(decoder_outputs[0]), out_past_key_values
-
-
-class OnnxT5(T5ForConditionalGeneration):
-    """creates a T5 model using onnx sessions (encode, decoder & init_decoder)"""
-
-    def __init__(self, model_or_model_path, onnx_model_sessions):
-        config = AutoConfig.from_pretrained(
-            model_or_model_path, use_auth_token=get_auth_token()
-        )
-        super().__init__(config)
-
-        # monkeypatch to work for MT5
-        if (
-            isinstance(model_or_model_path, str)
-            and "mt5" in model_or_model_path.lower()
-        ) or (
-            hasattr(model_or_model_path, "name_or_path")
-            and "mt5" in model_or_model_path.name_or_path
-        ):
-            self.model_type = "mt5"
-            self.config_class = MT5Config
-            self._keys_to_ignore_on_load_missing = [
-                r"encoder\.embed_tokens\.weight",
-            ]
-            self._keys_to_ignore_on_save = [
-                r"encoder\.embed_tokens\.weight",
-            ]
-
-        assert len(onnx_model_sessions) == 3, "all three models should be given"
-
-        encoder_sess, decoder_sess, decoder_sess_init = onnx_model_sessions
-
-        self.encoder = T5Encoder(encoder_sess)
-        self.decoder = T5Decoder(decoder_sess)
-        self.decoder_init = T5DecoderInit(decoder_sess_init)
-
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        decoder_input_ids=None,
-        decoder_attention_mask=None,
-        head_mask=None,
-        decoder_head_mask=None,
-        cross_attn_head_mask=None,
-        encoder_outputs=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        decoder_inputs_embeds=None,
-        labels=None,
-        use_cache=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        return_dict=None,
-    ):
-
-        if encoder_outputs is None:
-            # Convert encoder inputs in embeddings if needed
-            encoder_outputs = self.encoder(
-                input_ids=input_ids, attention_mask=attention_mask
-            )
-
-        encoder_hidden_states = encoder_outputs[0]
-
-        if past_key_values is not None:
-            if decoder_input_ids is not None:
-                decoder_input_ids = decoder_input_ids[:, -1:]
-            if decoder_inputs_embeds is not None:
-                decoder_inputs_embeds = decoder_inputs_embeds[:, -1:]
-
-        if past_key_values is None:
-
-            # runs only for the first time:
-            init_onnx_outputs = self.decoder_init(
-                decoder_input_ids, attention_mask, encoder_hidden_states
-            )
-
-            logits, past_key_values = init_onnx_outputs
-
-        else:
-
-            onnx_outputs = self.decoder(
-                decoder_input_ids,
-                attention_mask,
-                encoder_hidden_states,
-                past_key_values,
-            )
-
-            logits, past_key_values = onnx_outputs
-
-        return Seq2SeqLMOutput(logits=logits, past_key_values=past_key_values)
-
-
-def export_and_get_onnx_model(
-    model_or_model_path, custom_output_path=saved_models_path, quantized=True
-):
-    """
-                          Method for whole pipeline,
-    converts from pytorch to onnx --> quantizes model --> sets onnx runtime
-                --> builds whole onnx model with all sessions
-
-    """
-
-    # Step 1. convert huggingfaces t5 model to onnx
-    onnx_model_paths = generate_onnx_representation(
-        model_or_model_path, output_path=custom_output_path
-    )
-
-    if quantized:
-        # Step 2. (recommended) quantize the converted model for fast inference and to reduce model size.
-        quant_model_paths = quantize(onnx_model_paths)
-
-        # step 3. setup onnx runtime
-        print("Setting up onnx model...")
-        model_sessions = get_onnx_runtime_sessions(quant_model_paths)
-    else:
-        print("Setting up onnx model...")
-        model_sessions = get_onnx_runtime_sessions(onnx_model_paths)
-
-    # step 4. get the onnx model
-    model = OnnxT5(model_or_model_path, model_sessions)
-    print("Done!")
-
-    return model
-
-
-def get_onnx_model(model_name, onnx_models_path=saved_models_path, quantized=True):
-    """
-    method gets the onnx model, if already converted models exists
-    Example:
-    >> get_onnx_model(model_name="t5-finetuned", onnx_models_path="../models/onnx/quantized/")
-
-    """
-
-    encoder_path, decoder_path, init_decoder_path = get_model_paths(
-        model_name, Path(onnx_models_path), quantized
-    )
-
-    if quantized:
-        assert (
-            encoder_path.exists()
-            and decoder_path.exists()
-            and init_decoder_path.exists()
-        ), "quantized model don't exist in the model folder, first quantize the model!"
-    else:
-        assert (
-            encoder_path.exists()
-            and decoder_path.exists()
-            and init_decoder_path.exists()
-        ), "all or some models don't exists in the model folder, first convert the model! "
-
-    model_paths = encoder_path, decoder_path, init_decoder_path
-
-    model_sessions = get_onnx_runtime_sessions(model_paths)
-
-    model = OnnxT5(model_name, model_sessions)
-
-    return model
+from FastT5 import get_onnx_runtime_sessions, OnnxT5
 
 
 trained_model_path = './t5_squad_v1/'