mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-11 21:41:18 +00:00
Justin John
50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
import os
|
|
import yaml
|
|
import torch
|
|
from transformers import AlbertConfig, AlbertModel
|
|
|
|
|
|
class CustomAlbert(AlbertModel):
|
|
def forward(self, *args, **kwargs):
|
|
# Call the original forward method
|
|
outputs = super().forward(*args, **kwargs)
|
|
|
|
# Only return the last_hidden_state
|
|
return outputs.last_hidden_state
|
|
|
|
|
|
def load_plbert(log_dir):
|
|
config_path = os.path.join(log_dir, "config.yml")
|
|
plbert_config = yaml.safe_load(open(config_path))
|
|
|
|
albert_base_configuration = AlbertConfig(**plbert_config["model_params"])
|
|
bert = CustomAlbert(albert_base_configuration)
|
|
|
|
files = os.listdir(log_dir)
|
|
ckpts = []
|
|
for f in os.listdir(log_dir):
|
|
if f.startswith("step_"):
|
|
ckpts.append(f)
|
|
|
|
iters = [
|
|
int(f.split("_")[-1].split(".")[0])
|
|
for f in ckpts
|
|
if os.path.isfile(os.path.join(log_dir, f))
|
|
]
|
|
iters = sorted(iters)[-1]
|
|
|
|
checkpoint = torch.load(log_dir + "/step_" + str(iters) + ".t7", map_location="cpu")
|
|
state_dict = checkpoint["net"]
|
|
from collections import OrderedDict
|
|
|
|
new_state_dict = OrderedDict()
|
|
for k, v in state_dict.items():
|
|
name = k[7:] # remove `module.`
|
|
if name.startswith("encoder."):
|
|
name = name[8:] # remove `encoder.`
|
|
new_state_dict[name] = v
|
|
del new_state_dict["embeddings.position_ids"]
|
|
bert.load_state_dict(new_state_dict, strict=False)
|
|
|
|
return bert
|