tryolabs/transformers-optimization

1import time2 3import gradio as gr4import torch5from huggingface_hub import hf_hub_download6from onnxruntime import InferenceSession7from transformers import AutoModelForQuestionAnswering, AutoTokenizer8 9MAX_SEQUENCE_LENGTH = 51210 11models = {12    "Base model": "madlag/bert-large-uncased-whole-word-masking-finetuned-squadv2",13    "Pruned model": "madlag/bert-large-uncased-wwm-squadv2-x2.63-f82.6-d16-hybrid-v1",14    "Pruned ONNX Optimized FP16": "tryolabs/bert-large-uncased-wwm-squadv2-optimized-f16",15}16 17loaded_models = {18    "Pruned ONNX Optimized FP16": hf_hub_download(19        repo_id=models["Pruned ONNX Optimized FP16"], filename="model.onnx"20    ),21    "Base model": AutoModelForQuestionAnswering.from_pretrained(models["Base model"]),22    "Pruned model": AutoModelForQuestionAnswering.from_pretrained(23        models["Pruned model"]24    ),25}26 27 28def run_ort_inference(model_name, inputs):29    sess = InferenceSession(30        loaded_models[model_name], providers=["CPUExecutionProvider"]31    )32    start_time = time.time()33    output = sess.run(None, input_feed=inputs)34    end_time = time.time()35    return (output[0], output[1]), (end_time - start_time)36 37 38def run_normal_hf(model_name, inputs):39    start_time = time.time()40    output = loaded_models[model_name](**inputs).values()41    end_time = time.time()42    return output, (end_time - start_time)43 44 45def inference(model_name, context, question):46    tokenizer = AutoTokenizer.from_pretrained(models[model_name])47    if model_name == "Pruned ONNX Optimized FP16":48        inputs = dict(49            tokenizer(50                question, context, return_tensors="np", max_length=MAX_SEQUENCE_LENGTH51            )52        )53        output, inference_time = run_ort_inference(model_name, inputs)54        answer_start_scores, answer_end_scores = torch.tensor(output[0]), torch.tensor(55            output[1]56        )57    else:58        inputs = tokenizer(59            question, context, return_tensors="pt", max_length=MAX_SEQUENCE_LENGTH60        )61        output, inference_time = run_normal_hf(model_name, inputs)62        answer_start_scores, answer_end_scores = output63 64    input_ids = inputs["input_ids"].tolist()[0]65    answer_start = torch.argmax(answer_start_scores)66    answer_end = torch.argmax(answer_end_scores) + 167    answer = tokenizer.convert_tokens_to_string(68        tokenizer.convert_ids_to_tokens(input_ids[answer_start:answer_end])69    )70 71    return answer, f"{inference_time:.4f}s"72 73 74model_field = gr.Dropdown(75    choices=["Base model", "Pruned model", "Pruned ONNX Optimized FP16"],76    value="Pruned ONNX Optimized FP16",77    label="Model",78)79input_text_field = gr.Textbox(placeholder="Enter the text here", label="Text")80input_question_field = gr.Text(placeholder="Enter the question here", label="Question")81 82output_model = gr.Text(label="Model output")83output_inference_time = gr.Text(label="Inference time in seconds")84 85 86examples = [87    [88        "Pruned ONNX Optimized FP16",89        "The first little pig was very lazy. He didn't want to work at all and he built his house out of straw. The second little pig worked a little bit harder but he was somewhat lazy too and he built his house out of sticks. Then, they sang and danced and played together the rest of the day.",90        "Who worked a little bit harder?",91    ]92]93 94demo = gr.Interface(95    inference,96    inputs=[model_field, input_text_field, input_question_field],97    outputs=[output_model, output_inference_time],98    examples=examples,99)100 101demo.launch()102