marsggbo/t5-small_dff2048_dmodel32_token-pattern-predictor_switch128_xsum

1. Download the repo

import os
import torch

from glob import glob
from transformers import AutoModelForSeq2SeqLM, AutoConfig

model_name = 'marsggbo/t5-small_dff2048_dmodel32_token-pattern-predictor_switch128_xsum'
# ignore the mismatched size, because lm_head was modified
model = AutoModelForSeq2SeqLM.from_pretrained(
  model_name, ignore_mismatched_sizes=True, use_safetensors=False
)

2. Build the model

home_path = os.path.expanduser('~')
num_classes = 128 # switch128
ckpt_path = f"{home_path}/.cache/huggingface/hub/*{model_name.split('/')[-1]}/snapshots/*/*bin"
ckpt_path = glob(ckpt_path)[0]

model_config = AutoConfig.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_config(config=model_config)
model.lm_head = torch.nn.Linear(model.config.hidden_size, num_classes*6, bias=False)
model.load_state_dict(torch.load(ckpt_path, map_location='cpu'), strict=True)