Spaces:
Sleeping
Sleeping
File size: 5,585 Bytes
31f6bcb | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """Inference utilities for AspectBERT.
- load_model(): load tokenizer + DistilBERT backbone + classifier head from a
local checkpoint directory or a HuggingFace Hub repo (HF_MODEL_NAME).
- predict_aspect() / predict_all_aspects(): run sentiment prediction for one
or all 8 aspects on a single review.
- explain_with_lime(): word-importance explanation for a (review, aspect) pair.
"""
import argparse
import json
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import numpy as np
import torch
import torch.nn.functional as F
from transformers import DistilBertModel, DistilBertTokenizerFast
from constants import ASPECTS, ID2LABEL, MAX_LENGTH, MODEL_NAME, format_input # noqa: E402
from model import AspectBERT # noqa: E402
def _resolve_classifier_head(model_path):
"""Find classifier_head.pt either locally or on the HF Hub."""
local_path = os.path.join(model_path, "classifier_head.pt")
if os.path.exists(local_path):
return local_path
try:
from huggingface_hub import hf_hub_download
return hf_hub_download(repo_id=model_path, filename="classifier_head.pt")
except Exception:
return None
def load_model(model_path=None, device=None):
"""Load AspectBERT for inference.
`model_path` may be:
- None: read from HF_MODEL_NAME env var, falling back to the base
distilbert-base-uncased weights with an untrained head.
- a local checkpoint directory (as written by train.save_checkpoint)
- a HuggingFace Hub repo id containing the backbone, tokenizer, and
classifier_head.pt
"""
device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_path = model_path or os.environ.get("HF_MODEL_NAME") or MODEL_NAME
tokenizer = DistilBertTokenizerFast.from_pretrained(model_path)
model = AspectBERT(model_name=MODEL_NAME)
if model_path != MODEL_NAME:
try:
model.distilbert = DistilBertModel.from_pretrained(model_path)
except Exception as exc:
print(f"Warning: could not load fine-tuned backbone from "
f"'{model_path}' ({exc}); using base distilbert weights.")
classifier_head_path = _resolve_classifier_head(model_path)
if classifier_head_path:
state_dict = torch.load(classifier_head_path, map_location="cpu")
model.classifier.load_state_dict(state_dict)
else:
print(f"Warning: no classifier_head.pt found for '{model_path}'; "
"classification head is randomly initialized (untrained).")
model.to(device)
model.eval()
return model, tokenizer, device
@torch.no_grad()
def predict_aspect(model, tokenizer, device, text, aspect):
"""Predict the sentiment label + class probabilities for one aspect."""
inp = format_input(text, aspect)
enc = tokenizer(inp, truncation=True, padding="max_length",
max_length=MAX_LENGTH, return_tensors="pt")
enc = {k: v.to(device) for k, v in enc.items()}
logits = model(enc["input_ids"], enc["attention_mask"])
probs = F.softmax(logits, dim=-1).cpu().numpy()[0]
label_idx = int(np.argmax(probs))
return {
"label": ID2LABEL[label_idx],
"scores": {ID2LABEL[i]: float(probs[i]) for i in range(len(probs))},
}
def predict_all_aspects(model, tokenizer, device, text, aspects=None):
"""Predict sentiment for every aspect (default: all 8) on one review."""
aspects = aspects or ASPECTS
return {aspect: predict_aspect(model, tokenizer, device, text, aspect) for aspect in aspects}
def explain_with_lime(model, tokenizer, device, text, aspect, num_features=10, num_samples=200):
"""Return a LIME explanation object for the (text, aspect) prediction.
Requires the `lime` package.
"""
from lime.lime_text import LimeTextExplainer
class_names = [ID2LABEL[i] for i in range(len(ID2LABEL))]
explainer = LimeTextExplainer(class_names=class_names)
def predict_proba(texts):
out = []
for t in texts:
inp = format_input(t, aspect)
enc = tokenizer(inp, truncation=True, padding="max_length",
max_length=MAX_LENGTH, return_tensors="pt")
enc = {k: v.to(device) for k, v in enc.items()}
with torch.no_grad():
logits = model(enc["input_ids"], enc["attention_mask"])
probs = F.softmax(logits, dim=-1).cpu().numpy()
out.append(probs[0])
return np.array(out)
return explainer.explain_instance(
text, predict_proba, num_features=num_features,
num_samples=num_samples, labels=list(range(len(class_names))),
)
def main():
parser = argparse.ArgumentParser(description="Run AspectBERT inference on a review.")
parser.add_argument("text", help="Review text to analyze.")
parser.add_argument("--model_path", default=None,
help="Local checkpoint dir or HF Hub repo id "
"(defaults to HF_MODEL_NAME env var).")
parser.add_argument("--aspect", default=None,
help="Single aspect to predict; default = all 8 aspects.")
args = parser.parse_args()
model, tokenizer, device = load_model(args.model_path)
if args.aspect:
result = {args.aspect: predict_aspect(model, tokenizer, device, args.text, args.aspect)}
else:
result = predict_all_aspects(model, tokenizer, device, args.text)
print(json.dumps(result, indent=2))
if __name__ == "__main__":
main()
|