"""
Loading helper for the BetterLens dual-head model.

Usage (note: requires trust_remote_code=True):

    from loading_utils import load_dual_head
    model, tokenizer = load_dual_head("starmatrixtechnologies/betterlens-text-classifier")

    out = model(**tokenizer("Something is just generally wrong these days.",
                           return_tensors="pt", max_length=128,
                           padding="max_length", truncation=True))
    import torch
    probs = torch.softmax(out.sentiment_logits, dim=-1)
    print(probs[0])            # [positive, neutral, negative]
    print(out.vagueness_score) # [0..1]
"""

from transformers import AutoTokenizer, AutoModel


def load_dual_head(model_id="starmatrixtechnologies/betterlens-text-classifier"):
    tokenizer = AutoTokenizer.from_pretrained(
        model_id, trust_remote_code=True, use_fast=True
    )
    model = AutoModel.from_pretrained(
        model_id, trust_remote_code=True, torch_dtype="float32"
    )
    model.eval()
    return model, tokenizer


def predict(model, tokenizer, text, max_length=128):
    """Single-text convenience wrapper. Returns a dict with labels + scores."""
    import torch

    enc = tokenizer(
        text, return_tensors="pt", max_length=max_length,
        padding="max_length", truncation=True,
    )
    with torch.no_grad():
        out = model(**enc)
    probs = torch.softmax(out.sentiment_logits, dim=-1)[0]
    label = out.sentiment_labels[0].item()
    names = list(model.config.sentiment_label_names)
    return {
        "text": text,
        "sentiment": names[label],
        "sentiment_probs": {n: float(p) for n, p in zip(names, probs)},
        "vagueness": float(out.vagueness_score[0, 0]),
    }