File size: 1,479 Bytes
502d2a7 3df6849 502d2a7 3df6849 502d2a7 3df6849 502d2a7 3df6849 502d2a7 3df6849 502d2a7 3df6849 | 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 | import torch
from transformers import AutoModel, AutoTokenizer
def load_model(repo_id, device=None, revision=None):
device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
tokenizer = AutoTokenizer.from_pretrained(repo_id, revision=revision)
model = AutoModel.from_pretrained(
repo_id,
revision=revision,
trust_remote_code=True,
use_safetensors=True,
attn_implementation='sdpa',
).to(device).eval()
return model, tokenizer
@torch.inference_mode()
def predict(model, tokenizer, states, calibrated=True):
if isinstance(states, str):
states = [states]
if calibrated and model.config.temperatures is None:
raise RuntimeError('This model has no fitted calibration parameters')
device = next(model.parameters()).device
tokens = tokenizer(
states,
padding=True,
truncation=True,
max_length=model.config.max_length,
return_tensors='pt',
).to(device)
logits = model(**tokens).logits
temperatures = model.config.temperatures if calibrated else [1.0] * len(logits)
probabilities = [torch.softmax(head.float() / value, dim=-1).cpu()
for head, value in zip(logits, temperatures)]
return [
{
spec['name']: dict(zip(spec['options'], values[row].tolist()))
for spec, values in zip(model.config.decisions, probabilities)
}
for row in range(len(states))
]
|