Spaces:
Runtime error
Runtime error
File size: 1,985 Bytes
fc7db98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | """Load trained LCP checkpoints with backward compatibility."""
from __future__ import annotations
from pathlib import Path
import numpy as np
import torch
from transformers import AutoTokenizer
from two_head_model import TwoHeadModel, add_tgt_tokens
from utils import ENCODING_SPAN_MARK, MODELS, POOLING_CLS_CONCAT, POOLING_SPAN
def load_model_from_checkpoint(ckpt_path: str | Path, device: torch.device | None = None):
ckpt_path = Path(ckpt_path)
device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
model_key = ckpt.get("model_key", "deberta")
encoding = ckpt.get("encoding", ENCODING_SPAN_MARK)
pooling = ckpt.get("pooling", POOLING_CLS_CONCAT) # legacy checkpoints
use_linguistic = ckpt.get("use_linguistic_features", False)
level_only = ckpt.get("level_only", False)
tokenizer = AutoTokenizer.from_pretrained(MODELS[model_key])
tgt_id, tgt_end_id = add_tgt_tokens(tokenizer)
model = TwoHeadModel(
model_name=MODELS[model_key],
pooling_mode=pooling,
use_linguistic_features=use_linguistic,
)
model.encoder.resize_token_embeddings(len(tokenizer))
model.load_state_dict(ckpt["model_state"])
model.set_tgt_token_ids(ckpt.get("tgt_id", tgt_id), ckpt.get("tgt_end_id", tgt_end_id))
model.to(device)
model.eval()
feat_mean = np.array(ckpt["feat_mean"]) if ckpt.get("feat_mean") is not None else None
feat_std = np.array(ckpt["feat_std"]) if ckpt.get("feat_std") is not None else None
meta = {
"model_key": model_key,
"encoding": encoding,
"pooling": pooling,
"use_linguistic_features": use_linguistic,
"level_only": level_only,
"tgt_id": ckpt.get("tgt_id", tgt_id),
"tgt_end_id": ckpt.get("tgt_end_id", tgt_end_id),
"feat_mean": feat_mean,
"feat_std": feat_std,
}
return model, tokenizer, meta
|