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))
    ]