File size: 1,773 Bytes
1cbad0f | 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 | """
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]),
}
|