# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + colab={"base_uri": "https://localhost:8080/"} id="6rWvi3WhdQob" outputId="826b55c1-a654-4908-f9d1-b897285a782e"
# !pip install transformers

# + [markdown] id="wtjmOfq0d6db"
# # Benchmarking with pytorch

# + id="hpFRjm90dv77"
from transformers import GPT2LMHeadModel
model_name = 'gpt2'
model = GPT2LMHeadModel.from_pretrained(model_name)

# + colab={"base_uri": "https://localhost:8080/"} id="Rx8oXihpfPF4" outputId="3d25636f-db90-4871-ae6a-0d002f169e44"
# Move to gpu
import torch
device = torch.device('cuda')
model.to(device)
model.eval()
print('Assigned')

if (torch.cuda.is_available()):
  print("Cuda Avaliable")

# + [markdown] id="xKzd2Xs2ejLM"
# ## Greedy
#

# + id="eYgRqYRGltg7"
from functools import wraps
import time
import numpy as np

average_runs = 4

def timeit(my_func):
    @wraps(my_func)
    def timed(*args, **kw):
        time_list = []
        for i in range(3):
          tstart = time.time()
          output = my_func(*args, **kw)
          tend = time.time()
          time_taken = tend-tstart
          time_list.append((tend - tstart))
          
        return np.mean(time_list[1:]) # return last 3
    return timed


# + colab={"base_uri": "https://localhost:8080/"} id="bQCN4u0eeQqo" outputId="5f72606a-d5d6-4a10-cf4e-b9a04d738b32"
timing_results = {}
batch_size_list = [1, 3, 5, 10]
max_seq_length_list = [15, 25, 50, 100, 150]

@timeit
def call_greedy(input_tensor, length):
  result_greedy = model.generate(input_tensor, max_length=length, early_stopping=False)
  return result_greedy

minval = 1
maxval = 50256
input_tensor_length = 100

timing_holder = {}
for batch in batch_size_list:
  for sequence_length in max_seq_length_list:
    input_tensor = torch.randint(minval, maxval, size=(batch,input_tensor_length), device=device)
    # because we need to genrate that much tokens from input_tensor_length
    average_time_taken = call_greedy(input_tensor, input_tensor_length + sequence_length)
    #timing_holder['{}_{}'.format(batch, sequence_length)] = average_time_taken
    timing_holder[(batch, sequence_length)] = average_time_taken
    print("Done batch_size {} sequence length {}".format(batch, sequence_length))

    
timing_results['greedy'] = timing_holder

# + [markdown] id="VovpEoPhwq1u"
# # Beam

# + id="6S-hEzVzgDTc"
batch_size_list = [1, 3, 5, 10]
max_seq_length_list = [15, 25, 50, 100, 150]
beam_list = [1, 2, 5, 7, 10]

@timeit
def call_beam(input_tensor, length, beam):
  result = model.generate(input_tensor, max_length=length, num_beams=beam, early_stopping=False)
  return result

minval = 1
maxval = 50256
input_tensor_length = 100

timing_holder = {}
for batch in batch_size_list:
  for sequence_length in max_seq_length_list:
    for beam in beam_list:
      input_tensor = torch.randint(minval, maxval, size=(batch,input_tensor_length), device=device)
      # because we need to genrate that much tokens from input_tensor_length
      average_time_taken = call_beam(input_tensor, input_tensor_length + sequence_length, beam)
      #timing_holder['{}_{}'.format(batch, sequence_length)] = average_time_taken
      timing_holder[(batch, sequence_length, beam)] = average_time_taken
      print("Done batch_size {} sequence length {}".format(batch, sequence_length))

    
timing_results['beam'] = timing_holder

# + [markdown] id="ovZfdICfyO1V"
# ## Top K Top P

# + id="KPGmc21PitMu"
batch_size_list = [1, 3, 5, 10]
max_seq_length_list = [15, 25, 50, 100, 150]
num_return_sequence_list = [1, 2, 5, 7, 10]

@timeit
def call_top_k_top_p(input_tensor, length, return_sequence):
  result = model.generate(input_tensor, max_length=length, num_return_sequences=return_sequence, top_p=0.75, 
    top_k=25, do_sample=True, early_stopping=False)
  return result

minval = 1
maxval = 50256
input_tensor_length = 100

timing_holder = {}
for batch in batch_size_list:
  for sequence_length in max_seq_length_list:
    for num_return_sequence in num_return_sequence_list:
      input_tensor = torch.randint(minval, maxval, size=(batch,input_tensor_length), device=device)
      # because we need to genrate that much tokens from input_tensor_length
      average_time_taken = call_top_k_top_p(input_tensor, input_tensor_length + sequence_length, num_return_sequence)
      #timing_holder['{}_{}'.format(batch, sequence_length)] = average_time_taken
      timing_holder[(batch, sequence_length, num_return_sequence)] = average_time_taken
      print("Done batch_size {} sequence length {}".format(batch, sequence_length))

    
timing_results['top_k_top_p'] = timing_holder

import pickle
with open("pytorch_gpt2_benchmark.pkl", "wb") as f:
    pickle.dump(timing_results, f)

# + id="ANmPzVqciu_E"


StarcoderdataJupyterScriptsDedupFiltered