class CodeGenTemplate:

  @classmethod
  def autogen_head(cls):
    return '''# Autogenerated by onnx-pytorch.
'''

  @classmethod
  def imports(cls):
    return '''import glob
import os

import numpy as np
import torch
import torch.nn as nn
'''

  @classmethod
  def model(cls, model_init, model_forward, test_run_model):
    return f'''{cls.autogen_head()}
{cls.imports()}


class Model(nn.Module):
  def __init__(self):
    super(Model, self).__init__()
    self.__variables = {{
        os.path.basename(b)[:-4]: torch.from_numpy(np.load(b))
        for b in glob.glob(
            os.path.join(os.path.dirname(__file__), "variables", "*.npy"))
        if os.path.isfile(b)
    }}
    {model_init}

  def forward(self, *inputs):
    {model_forward}


{test_run_model}
'''
