Views
No views yet
1import os
2import torch
3
4from glob import glob
5from transformers import AutoModelForSeq2SeqLM, AutoConfig
6
7model_name = 'marsggbo/t5-small_dff2048_dmodel32_token-pattern-predictor_switch32_wmt16'
8# ignore the mismatched size, because lm_head was modified
9model = AutoModelForSeq2SeqLM.from_pretrained(
10 model_name, ignore_mismatched_sizes=True, use_safetensors=False
11)1home_path = os.path.expanduser('~')
2num_classes = 32 # switch32
3ckpt_path = f"{home_path}/.cache/huggingface/hub/*{model_name.split('/')[-1]}/snapshots/*/*bin"
4ckpt_path = glob(ckpt_path)[0]
5
6model_config = AutoConfig.from_pretrained(model_name)
7model = AutoModelForSeq2SeqLM.from_config(config=model_config)
8model.lm_head = torch.nn.Linear(model.config.hidden_size, num_classes*6, bias=False)
9model.load_state_dict(torch.load(ckpt_path, map_location='cpu'), strict=True)