Skip to content

Commit fa5cfdc

Browse files
feat(backends): Add native Hugging Face transformers backend resolver (#168)
* feat: add transformers resolver * tests: add tests for huggingface resolver handling of id2label and y_prob * fix: update installation instructions for huggingface resolver to use `trustlens[nlp]` * enhance: centralize ordered label derivation in a shared helper _ordered_labels_from_id2label * style: resolve lints and type checks * style: improve type hints and label handling in id2label functions --------- Co-authored-by: Shahid Ul Islam <65996596+Khanz9664@users.noreply.github.com>
1 parent ef29d43 commit fa5cfdc

4 files changed

Lines changed: 459 additions & 0 deletions

File tree

‎pyproject.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,12 +68,16 @@ torch = [
6868
"torch>=2.0",
6969
"torchvision>=0.15",
7070
]
71+
nlp = [
72+
"transformers>=4.30",
73+
]
7174
full = [
7275
"xgboost>=1.7",
7376
"lightgbm>=3.3",
7477
"catboost>=1.0",
7578
"torch>=2.0",
7679
"torchvision>=0.15",
80+
"transformers>=4.30",
7781
"shap>=0.42",
7882
"captum>=0.6",
7983
]

‎tests/test_backend_huggingface.py‎

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
from types import SimpleNamespace
2+
3+
import numpy as np
4+
import pytest
5+
6+
from trustlens import TrustReport, analyze
7+
from trustlens.backends.registry import detect_framework, get_resolver
8+
from trustlens.backends.types import UnsupportedModelError
9+
10+
# We only run these tests if transformers is available
11+
transformers = pytest.importorskip("transformers")
12+
13+
14+
def _make_text_classification_pipeline(id2label, outputs):
15+
"""
16+
Build a lightweight double of a `TextClassificationPipeline` that subclasses
17+
the real class (so `type(model).__module__` starts with 'transformers' and
18+
`isinstance(model, TextClassificationPipeline)` is True) without invoking the
19+
real `__init__`, which requires a loaded model/tokenizer.
20+
"""
21+
from transformers import TextClassificationPipeline
22+
23+
class MockTextClassificationPipeline(TextClassificationPipeline):
24+
def __init__(self):
25+
self.task = "text-classification"
26+
self.model = SimpleNamespace(config=SimpleNamespace(id2label=id2label))
27+
28+
def __call__(self, X, top_k=None):
29+
return outputs
30+
31+
MockTextClassificationPipeline.__module__ = TextClassificationPipeline.__module__
32+
return MockTextClassificationPipeline()
33+
34+
35+
def _make_non_classification_pipeline():
36+
"""
37+
Build a double of a `TokenClassificationPipeline` (NER) — a real transformers
38+
pipeline class (so it's still detected as 'huggingface' via module prefix) but
39+
not a text-classification pipeline, so the resolver should reject it.
40+
"""
41+
from transformers import TokenClassificationPipeline
42+
43+
class MockTokenClassificationPipeline(TokenClassificationPipeline):
44+
def __init__(self):
45+
self.task = "ner"
46+
self.model = SimpleNamespace(config=SimpleNamespace())
47+
48+
def __call__(self, X, **kwargs):
49+
return [[] for _ in X]
50+
51+
MockTokenClassificationPipeline.__module__ = TokenClassificationPipeline.__module__
52+
return MockTokenClassificationPipeline()
53+
54+
55+
def test_huggingface_detection():
56+
model = _make_text_classification_pipeline(
57+
id2label={0: "NEGATIVE", 1: "POSITIVE"},
58+
outputs=[[{"label": "NEGATIVE", "score": 0.5}, {"label": "POSITIVE", "score": 0.5}]],
59+
)
60+
assert detect_framework(model) == "huggingface"
61+
62+
63+
def test_huggingface_resolver_basic():
64+
model = _make_text_classification_pipeline(
65+
id2label={0: "NEGATIVE", 1: "POSITIVE"},
66+
outputs=[
67+
[{"label": "POSITIVE", "score": 0.9}, {"label": "NEGATIVE", "score": 0.1}],
68+
[{"label": "NEGATIVE", "score": 0.7}, {"label": "POSITIVE", "score": 0.3}],
69+
],
70+
)
71+
X = ["great movie", "bad movie"]
72+
73+
resolver = get_resolver(model)
74+
bundle = resolver(model, X)
75+
76+
assert bundle.framework == "huggingface"
77+
assert bundle.y_pred.shape == (2,)
78+
assert bundle.y_prob.shape == (2, 2)
79+
assert bundle.metadata["resolver"] == "huggingface"
80+
assert list(bundle.class_labels) == ["NEGATIVE", "POSITIVE"]
81+
assert list(bundle.y_pred) == ["POSITIVE", "NEGATIVE"]
82+
83+
84+
def test_huggingface_integration_analyze():
85+
id2label = {0: "NEGATIVE", 1: "NEUTRAL", 2: "POSITIVE"}
86+
model = _make_text_classification_pipeline(
87+
id2label=id2label,
88+
outputs=[
89+
[
90+
{"label": "POSITIVE", "score": 0.7},
91+
{"label": "NEUTRAL", "score": 0.2},
92+
{"label": "NEGATIVE", "score": 0.1},
93+
],
94+
[
95+
{"label": "NEGATIVE", "score": 0.6},
96+
{"label": "NEUTRAL", "score": 0.3},
97+
{"label": "POSITIVE", "score": 0.1},
98+
],
99+
[
100+
{"label": "NEUTRAL", "score": 0.5},
101+
{"label": "POSITIVE", "score": 0.3},
102+
{"label": "NEGATIVE", "score": 0.2},
103+
],
104+
[
105+
{"label": "NEGATIVE", "score": 0.5},
106+
{"label": "POSITIVE", "score": 0.3},
107+
{"label": "NEUTRAL", "score": 0.2},
108+
],
109+
],
110+
)
111+
X = ["great movie", "terrible movie", "it was fine", "not for me"]
112+
y = np.array(["POSITIVE", "NEGATIVE", "NEUTRAL", "NEGATIVE"])
113+
114+
report = analyze(model, X, y, verbose=False)
115+
116+
assert isinstance(report, TrustReport)
117+
assert report.metadata["framework"] == "huggingface"
118+
assert "huggingface" in report.metadata["backend"]["resolver"]
119+
120+
121+
def test_huggingface_manual_override():
122+
id2label = {0: "NEGATIVE", 1: "NEUTRAL", 2: "POSITIVE"}
123+
model = _make_text_classification_pipeline(
124+
id2label=id2label,
125+
outputs=[
126+
[
127+
{"label": "POSITIVE", "score": 0.7},
128+
{"label": "NEUTRAL", "score": 0.2},
129+
{"label": "NEGATIVE", "score": 0.1},
130+
],
131+
[
132+
{"label": "NEGATIVE", "score": 0.6},
133+
{"label": "NEUTRAL", "score": 0.3},
134+
{"label": "POSITIVE", "score": 0.1},
135+
],
136+
],
137+
)
138+
X = ["great movie", "terrible movie"]
139+
y = np.array(["POSITIVE", "NEGATIVE"])
140+
141+
custom_preds = np.array(["POSITIVE", "NEUTRAL"])
142+
report = analyze(model, X, y, y_pred=custom_preds, verbose=False)
143+
144+
assert np.array_equal(report.y_pred, custom_preds)
145+
assert report.metadata["framework"] == "huggingface"
146+
147+
148+
def test_huggingface_non_classification_pipeline_rejection():
149+
model = _make_non_classification_pipeline()
150+
X = ["some long document to summarize"]
151+
y = np.array([0])
152+
153+
with pytest.raises(UnsupportedModelError):
154+
analyze(model, X, y, verbose=False)
155+
156+
157+
def test_huggingface_resolver_no_id2label_fallback():
158+
model = _make_text_classification_pipeline(
159+
id2label=None,
160+
outputs=[
161+
[{"label": "joy", "score": 0.6}, {"label": "anger", "score": 0.4}],
162+
[{"label": "anger", "score": 0.2}, {"label": "joy", "score": 0.8}],
163+
[{"label": "joy", "score": 0.3}, {"label": "anger", "score": 0.7}],
164+
],
165+
)
166+
X = ["a", "b", "c"]
167+
168+
resolver = get_resolver(model)
169+
bundle = resolver(model, X)
170+
171+
assert bundle.framework == "huggingface"
172+
assert bundle.y_prob.shape == (3, 2)
173+
# Column order comes from the first sample: joy=col0, anger=col1.
174+
assert list(bundle.class_labels) == ["joy", "anger"]
175+
np.testing.assert_allclose(bundle.y_prob, [[0.6, 0.4], [0.8, 0.2], [0.3, 0.7]])
176+
assert list(bundle.y_pred) == ["joy", "joy", "anger"]
177+
178+
179+
def test_huggingface_resolver_no_id2label_inconsistent_labels_raises():
180+
model = _make_text_classification_pipeline(
181+
id2label=None,
182+
outputs=[
183+
[{"label": "joy", "score": 0.6}, {"label": "anger", "score": 0.4}],
184+
[{"label": "sadness", "score": 0.5}, {"label": "anger", "score": 0.5}],
185+
],
186+
)
187+
X = ["a", "b"]
188+
189+
resolver = get_resolver(model)
190+
with pytest.raises(ValueError, match="not present in the canonical label set"):
191+
resolver(model, X)
192+
193+
194+
def test_huggingface_resolver_explicit_y_prob_2d():
195+
class ExplodingPipelineError(Exception):
196+
pass
197+
198+
def _boom(*_args, **_kwargs):
199+
raise ExplodingPipelineError("model(...) should not be called when y_prob is supplied")
200+
201+
model = _make_text_classification_pipeline(
202+
id2label={0: "NEGATIVE", 1: "POSITIVE"},
203+
outputs=[], # unused; __call__ is overridden below to assert it's never hit
204+
)
205+
model.__class__.__call__ = _boom
206+
207+
explicit_y_prob = np.array([[0.2, 0.8], [0.9, 0.1]])
208+
resolver = get_resolver(model)
209+
bundle = resolver(model, ["a", "b"], y_prob=explicit_y_prob)
210+
211+
assert bundle.framework == "huggingface"
212+
np.testing.assert_array_equal(bundle.y_prob, explicit_y_prob)
213+
assert list(bundle.class_labels) == ["NEGATIVE", "POSITIVE"]
214+
assert list(bundle.y_pred) == ["POSITIVE", "NEGATIVE"]
215+
216+
217+
def test_huggingface_resolver_explicit_y_prob_1d_normalized():
218+
model = _make_text_classification_pipeline(
219+
id2label={0: "NEGATIVE", 1: "POSITIVE"},
220+
outputs=[],
221+
)
222+
223+
def _boom(*_args, **_kwargs):
224+
raise AssertionError("model(...) should not be called when y_prob is supplied")
225+
226+
model.__class__.__call__ = _boom
227+
228+
explicit_y_prob_1d = np.array([0.8, 0.2])
229+
resolver = get_resolver(model)
230+
bundle = resolver(model, ["a", "b"], y_prob=explicit_y_prob_1d)
231+
232+
assert bundle.y_prob.shape == (2, 2)
233+
np.testing.assert_allclose(bundle.y_prob, [[0.2, 0.8], [0.8, 0.2]])
234+
assert list(bundle.y_pred) == ["POSITIVE", "NEGATIVE"]

0 commit comments

Comments
 (0)