diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0432a655..fbc40cea 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -32,25 +32,25 @@ jobs: name: Build and test (Python ${{ matrix.python-version }}, Node ${{ matrix.node-version }}) strategy: matrix: - node-version: [18] - python-version: ["3.9", "3.10", "3.11"] + node-version: [22] + python-version: ["3.12", "3.13", "3.14"] defaults: run: shell: bash -l {0} runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v2 + uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install LIT package with testing dependencies - run: python -m pip install -e '.[test]' + run: python -m pip install -e '.[examples-common,examples-discriminative-ai,examples-generative-ai,umap,test]' - name: Test LIT run: pytest -v - name: Setup Node ${{ matrix.node-version }} - uses: actions/setup-node@v2 + uses: actions/setup-node@v4 with: node-version: ${{ matrix.node-version }} - name: Build LIT client diff --git a/Dockerfile b/Dockerfile index 5ca68273..7de0f8dc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -17,7 +17,7 @@ # ---- LIT Base Container ---- -FROM python:3.11-slim AS lit-nlp-base +FROM python:3.13-slim AS lit-nlp-base # Update Ubuntu packages and install basic utils RUN apt-get update @@ -35,7 +35,7 @@ COPY ./lit_nlp/examples/gunicorn_config.py ./ FROM lit-nlp-base AS lit-nlp-prod -RUN python -m pip install 'lit-nlp[examples-discriminative-ai]' +RUN python -m pip install 'lit-nlp[examples-common,examples-discriminative-ai]' WORKDIR $APP_HOME ENTRYPOINT ["gunicorn", "--config=gunicorn_config.py"] diff --git a/README.md b/README.md index 13e74503..ae0d3665 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ Features include: * **Highly extensible** to new model types, including classification, regression, span labeling, seq2seq, and language modeling. Supports multi-head models and multiple input features out of the box. -* **Framework-agnostic** and compatible with TensorFlow, PyTorch, and more. +* **Framework-agnostic** and compatible with PyTorch, scikit-learn, and more. LIT has a [website](https://pair-code.github.io/lit) with live demos, tutorials, a setup guide and more. @@ -67,13 +67,13 @@ appropriate optional dependencies. ```sh # To install dependencies for the discriminative AI examples (GLUE, Penguin) -pip install 'lit-nlp[examples-discriminative-ai]' +pip install 'lit-nlp[examples-common,examples-discriminative-ai]' # To install dependencies for the generative AI examples (Prompt Debugging) -pip install 'lit-nlp[examples-generative-ai]' +pip install 'lit-nlp[examples-common,examples-generative-ai]' # To install dependencies for all examples plus the test suite -pip install 'lit-nlp[test]' +pip install 'lit-nlp[examples-common,examples-discriminative-ai,examples-generative-ai,umap,test]' ``` ### Install from source @@ -85,8 +85,8 @@ git clone https://github.com/PAIR-code/lit.git cd lit ``` -Note: be sure you are running Python 3.9+. If you have a different version on -your system, use the `conda` instructions below to set up a Python 3.9 +Note: be sure you are running Python 3.12+. If you have a different version on +your system, use the `conda` instructions below to set up a Python 3.12 environment. Set up a Python environment with `venv` (or your preferred environment manager). @@ -97,7 +97,7 @@ other optional dependency possibilities in the install with pip section. ```sh python -m venv .venv source .venv/bin/activate -python -m pip install -e '.[test]' +python -m pip install -e '.[examples-common,examples-discriminative-ai,examples-generative-ai,umap,test]' ``` The LIT repo does not include a distributable version of the LIT app. You must diff --git a/RELEASE.md b/RELEASE.md index 09591ed8..5422cef9 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,6 +1,35 @@ # Learning Interpretability Tool Release Notes +## Release 1.4 + +This release completes the migration of LIT's bundled models, examples, and +dependencies from TensorFlow to PyTorch. TensorFlow is no longer a required or +supported dependency of any LIT example. + +### Breaking changes +* GLUE models (`examples/glue`) are now implemented with PyTorch + (`AutoModelForSequenceClassification`); the Keras `train()` history format + and callback arguments are replaced by a plain dict and an `on_epoch_end` + callback. +* GLUE, TyDi, and penguin datasets now load from HuggingFace `datasets` (or + CSV) instead of TensorFlow Datasets. +* The penguin demo model is now a scikit-learn classifier trained on the + dataset at init; the hosted `penguin.h5` asset is no longer used. +* Prompt debugging now supports HuggingFace Transformers on PyTorch only; the + KerasNLP wrapper (`keras_lms.py`) and the `dl_framework` option have been + removed. +* TyDi QA demo models now run with PyTorch instead of Flax. +* Optional extras are now additive and disjoint; combine them when installing, + e.g. `pip install 'lit-nlp[examples-common,examples-discriminative-ai]'`. + The old self-contained extras are replaced by + `examples-common` + a task-specific extra. + +### Dependencies +* Removed `tensorflow`, `tf-keras`, `keras-nlp`, `tensorflow-text`, and + `tensorflow-datasets`; added `torch`, `datasets`, and `huggingface-hub`. + + ## Release 1.3.1 This is a minor update to fix issues with running the [LIT Gemma Colab](https://colab.sandbox.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/lit_gemma.ipynb). diff --git a/lit_nlp/api/model.py b/lit_nlp/api/model.py index 79558fce..9ce23cdd 100644 --- a/lit_nlp/api/model.py +++ b/lit_nlp/api/model.py @@ -49,11 +49,11 @@ def maybe_copy_np(arr): # If this is not a view of another array. if arr.base is None: return arr - # Tensorflow provides a bridge to share memory between tensorflow and numpy - # arrays. This looks like a view into an array but the base is a - # tensorflow_wrapper not an array, so the view heuristics below don't work. We - # can check for this case by checking is arr.base has the ndim attribute. - # https://github.com/tensorflow/tensorflow/blob/6ed79e8429730c33dc894175da7a1849a8e3e57f/tensorflow/python/lib/core/ndarray_tensor_bridge.cc#L90 + # Frameworks like PyTorch provide bridges to share memory + # between their tensors and numpy arrays. This looks like a view into an array + # but the base is a tensor_wrapper not an array, so the view heuristics below + # don't work. We can check for this case by checking is arr.base has the ndim + # attribute. if not hasattr(arr.base, 'ndim'): return np.copy(arr) # Heuristic to check if we should 'detach' this array from the parent blob. diff --git a/lit_nlp/components/ablation_flip_int_test.py b/lit_nlp/components/ablation_flip_int_test.py index 10cf9f9e..299e26a5 100644 --- a/lit_nlp/components/ablation_flip_int_test.py +++ b/lit_nlp/components/ablation_flip_int_test.py @@ -15,6 +15,7 @@ """Tests for lit_nlp.components.ablation_flip.""" from collections.abc import Iterable, Iterator +import os from absl.testing import absltest from lit_nlp.api import types @@ -23,15 +24,11 @@ import numpy as np -from lit_nlp.lib import file_cache -BERT_TINY_PATH = file_cache.cached_path( - 'https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz', # pylint: disable=line-too-long - extract_compressed_file=True, -) -STSB_PATH = file_cache.cached_path( - 'https://storage.googleapis.com/what-if-tool-resources/lit-models/stsb_tiny.tar.gz', # pylint: disable=line-too-long - extract_compressed_file=True, +_TESTDATA_PATH = os.path.join( + os.path.dirname(glue_models.__file__), 'testdata' ) +BERT_TINY_PATH = os.path.join(_TESTDATA_PATH, 'sst2_tiny') +STSB_PATH = os.path.join(_TESTDATA_PATH, 'stsb_tiny') class SST2ModelNonRequiredField(glue_models.SST2Model): @@ -73,7 +70,7 @@ def setUp(self): # Classification model with a counter to count number of predict calls. # TODO(ataly): Consider setting up a Mock object to count number of - # predict calls. + # predict calls. self.classification_model_with_predict_counter = ( SST2ModelWithPredictCounter(BERT_TINY_PATH)) diff --git a/lit_nlp/components/citrus/utils_test.py b/lit_nlp/components/citrus/utils_test.py index db67268b..3349cfcb 100644 --- a/lit_nlp/components/citrus/utils_test.py +++ b/lit_nlp/components/citrus/utils_test.py @@ -17,10 +17,9 @@ from absl.testing import parameterized from lit_nlp.components.citrus import utils import numpy as np -import tensorflow.compat.v2 as tf -class UtilsTest(parameterized.TestCase, tf.test.TestCase): +class UtilsTest(parameterized.TestCase): def setUp(self): """Resets random seed for each test.""" @@ -54,7 +53,7 @@ def test_normalize_scores(self, scores, make_positive): """Check if the scores sum to 1 after taking their absolute values.""" original_min = np.min(scores) scores = utils.normalize_scores(scores, make_positive=make_positive) - self.assertAllClose(1.0, np.abs(scores).sum(-1)) + np.testing.assert_allclose(1.0, np.abs(scores).sum(-1)) if not make_positive: # Keep the sign of originally negative values. self.assertLessEqual(np.min(scores), np.max([0.0, original_min])) diff --git a/lit_nlp/components/hotflip_int_test.py b/lit_nlp/components/hotflip_int_test.py index c5d7758b..e0b13d57 100644 --- a/lit_nlp/components/hotflip_int_test.py +++ b/lit_nlp/components/hotflip_int_test.py @@ -14,6 +14,8 @@ # ============================================================================== """Tests for lit_nlp.components.hotflip.""" +import os + from absl.testing import absltest from absl.testing import parameterized from lit_nlp.components import hotflip @@ -21,15 +23,11 @@ from lit_nlp.examples.glue import models as glue_models import numpy as np -from lit_nlp.lib import file_cache -BERT_TINY_PATH = file_cache.cached_path( - 'https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz', # pylint: disable=line-too-long - extract_compressed_file=True, -) -STSB_PATH = file_cache.cached_path( - 'https://storage.googleapis.com/what-if-tool-resources/lit-models/stsb_tiny.tar.gz', # pylint: disable=line-too-long - extract_compressed_file=True, +_TESTDATA_PATH = os.path.join( + os.path.dirname(glue_models.__file__), 'testdata' ) +BERT_TINY_PATH = os.path.join(_TESTDATA_PATH, 'sst2_tiny') +STSB_PATH = os.path.join(_TESTDATA_PATH, 'stsb_tiny') _CONFIG_CLASSIFICATION = { hotflip.FIELDS_TO_HOTFLIP_KEY: ['tokens_sentence'], diff --git a/lit_nlp/components/image_gradient_maps.py b/lit_nlp/components/image_gradient_maps.py index 1e6e6199..dbf173d9 100644 --- a/lit_nlp/components/image_gradient_maps.py +++ b/lit_nlp/components/image_gradient_maps.py @@ -71,6 +71,38 @@ _SUPPORTED_PRED_TYPES = (types.MulticlassPreds, types.RegressionScore) +def _patch_saliency_numpy_compat() -> None: + """Makes saliency 0.2.1 work with NumPy >= 2.0. + + saliency.core.guided_ig calls np.quantile(..., interpolation='lower'), but + the `interpolation` keyword was removed in NumPy 2.0 in favor of `method`. + Swap the numpy reference used by that module for a proxy whose quantile + accepts the legacy keyword. Can be removed once a saliency release includes + the upstream fix. + """ + import types as stdlib_types + + from saliency.core import guided_ig + + if getattr(guided_ig, '_lit_numpy_compat', False): + return + original_quantile = np.quantile + + def _quantile(a, *args, **kwargs): + if 'interpolation' in kwargs: + kwargs['method'] = kwargs.pop('interpolation') + return original_quantile(a, *args, **kwargs) + + numpy_compat = stdlib_types.ModuleType('numpy') + numpy_compat.__dict__.update(np.__dict__) + numpy_compat.quantile = _quantile + guided_ig.np = numpy_compat + guided_ig._lit_numpy_compat = True + + +_patch_saliency_numpy_compat() + + class SupportedFields(NamedTuple): """The collection of field names that are required to calculate saliency.""" grad_field_key: str diff --git a/lit_nlp/components/tcav_int_test.py b/lit_nlp/components/tcav_int_test.py index 041a6a84..4a1af9aa 100644 --- a/lit_nlp/components/tcav_int_test.py +++ b/lit_nlp/components/tcav_int_test.py @@ -14,6 +14,7 @@ # ============================================================================== """Tests for lit_nlp.components.gradient_maps.""" +import os import random from absl.testing import absltest @@ -23,12 +24,12 @@ from lit_nlp.examples.glue import models as glue_models from lit_nlp.lib import caching # for hash id fn from lit_nlp.lib import testing_utils +import numpy as np -from lit_nlp.lib import file_cache -BERT_TINY_PATH = file_cache.cached_path( - 'https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz', # pylint: disable=line-too-long - extract_compressed_file=True, +_TESTDATA_PATH = os.path.join( + os.path.dirname(glue_models.__file__), 'testdata' ) +BERT_TINY_PATH = os.path.join(_TESTDATA_PATH, 'sst2_tiny') _ALPHABET_EXAMPLES = [ {'sentence': 'a', '_id': 'a'}, @@ -42,6 +43,24 @@ {'sentence': 'i', '_id': 'i'}, ] +def _unit_scale_cav(result): + """Normalizes the trained CAV and dot products to unit CAV scale. + + The magnitude of scikit-learn LinearSVC coefficients (and hence the CAV) + varies across library versions without changing the CAV direction, which is + what TCAV scores depend on. Normalizing keeps the golden values stable + across scikit-learn upgrades. + """ + result = dict(result) + inner = dict(result['result']) + cav = np.asarray(inner['cav'], dtype=np.float64) + scale = np.linalg.norm(cav) + inner['cav'] = (cav / scale).tolist() + inner['dot_prods'] = [d / scale for d in inner['dot_prods']] + result['result'] = inner + return result + + _EMOTION_EXAMPLES = [ {'sentence': 'happy', '_id': 'happy'}, # 0 {'sentence': 'sad', '_id': 'sad'}, # 1 @@ -111,56 +130,159 @@ def test_tcav(self): ) self.assertLen(result, 1) - expected = { - 'p_val': 0.13311, - 'random_mean': 0.56667, - 'result': { - 'score': - 0.33333, - 'cos_sim': [ - 0.088691, -0.12179, 0.16013, 0.24840, -0.09793, 0.05166, - -0.21578, -0.06560, -0.14759 - ], - 'dot_prods': [ - 189.085096, -266.36317, 344.350498, 547.144949, -211.663965, - 112.502439, -472.72066, -144.529598, -323.31888 - ], - 'accuracy': - 0.66667, - 'cav': [[ - 0.91326976, -13.67988429, -1.33965892, 12.23702997, 4.83590182, - -20.6276055, 19.07518597, 4.9641991, -5.52570226, 10.37834134, - 0.53685622, -5.21721739, 12.00234068, -26.36114057, 6.20088358, - 4.76729567, 26.4359945, 16.30961659, 14.70620605, -13.47771528, - 5.29218365, 0.87290488, 10.19441762, 2.96687215, -0.70745918, - -12.201803, -10.65010904, -5.90814342, -25.17510006, - -5.90629019, 26.53638293, 27.44054429, -11.75430314, - -3.72779609, 10.28197421, -11.58444132, -18.09351946, - -22.09520524, -30.04023056, -16.7551763, -12.34913637, - 30.27897114, -13.98790656, 0.65758253, 3.75770261, 4.81132118, - -11.18005293, 5.85445303, 18.88336475, -7.34885733, - -35.76094848, -3.39495953, 5.52774132, -8.38126488, 10.07413613, - -1.96825956, 17.97041089, -2.47085774, 19.5700424, -49.46295186, - 0.12541183, -6.95592842, -0.33562953, 12.9269603, -13.6288284, - -9.51211543, 21.22778867, -2.81344371, -9.66434107, - -15.41551695, -30.15406401, 14.65903841, 2.51729041, - 17.70711379, 13.21615045, 12.55899318, 11.45749114, - -25.67659992, 13.00876054, -1.52499005, 27.45026658, - -4.36110401, 10.01664277, -11.24470769, -6.79411522, - -1.67106503, -2.59389537, 11.72310319, -0.84126818, 2.03886137, - -11.25047383, -8.60679631, -23.1676229, 22.83532544, - -25.2657054, -15.49795527, -1.62890474, -15.49504251, - 27.26973702, -8.00652979, 6.87541861, -3.61878753, 0.82889822, - -2.88891667, 6.13730358, 12.55884424, -2.24121286, 45.90285087, - 34.43108722, -13.32567113, 19.0988537, 2.16242269, 17.45654791, - 18.17472208, 18.28023357, -20.34869744, -11.21275755, 6.0583063, - -10.38857432, -7.30056744, -8.85395997, -17.19779724, - -8.12087822, 2.46058364, 29.35061395, -6.53820405, 10.3522653, - 8.54478485 - ]] - } - } - testing_utils.assert_deep_almost_equal(self, expected, result[0]) + expected = {'p_val': 0.10634367, + 'random_mean': 0.54444444, + 'result': {'accuracy': 0.66666667, + 'cav': [[0.0052003721, + -0.07789658, + -0.0076281652, + 0.069680482, + 0.02753685, + -0.11745859, + 0.10861865, + 0.028267169, + -0.031464793, + 0.059096765, + 0.0030570414, + -0.029708052, + 0.06834428, + -0.15010664, + 0.03530933, + 0.027146045, + 0.15053286, + 0.092870869, + 0.083740711, + -0.076745257, + 0.030134914, + 0.0049704607, + 0.058049336, + 0.016893858, + -0.0040285308, + -0.06947986, + -0.060644325, + -0.03364237, + -0.14335285, + -0.033631887, + 0.15110439, + 0.156253, + -0.066931739, + -0.02122692, + 0.058548115, + -0.06596458, + -0.10302878, + -0.12581527, + -0.17105629, + -0.095408045, + -0.070318855, + 0.17241567, + -0.079650335, + 0.0037445326, + 0.021397213, + 0.027396964, + -0.063661948, + 0.033336654, + 0.10752612, + -0.041846175, + -0.20363152, + -0.019331753, + 0.031476293, + -0.047724877, + 0.057364281, + -0.011207522, + 0.10232788, + -0.014069673, + 0.11143632, + -0.28165388, + 0.00071425573, + -0.039608657, + -0.0019110381, + 0.073609203, + -0.077605829, + -0.054164339, + 0.12087602, + -0.016020484, + -0.055031046, + -0.087779574, + -0.17170441, + 0.083472103, + 0.014334007, + 0.1008285, + 0.075256057, + 0.071513988, + 0.06524168, + -0.14620869, + 0.074075021, + -0.008683526, + 0.15630831, + -0.024833282, + 0.05703707, + -0.064029969, + -0.038687266, + -0.0095154466, + -0.014770232, + 0.066754095, + -0.0047905217, + 0.011609488, + -0.064062886, + -0.049009111, + -0.13192189, + 0.13002953, + -0.14386889, + -0.088248827, + -0.0092752138, + -0.088232592, + 0.1552804, + -0.045591053, + 0.039150309, + -0.020606294, + 0.0047199279, + -0.016450174, + 0.03494728, + 0.071513101, + -0.012762167, + 0.26138169, + 0.19605877, + -0.075879507, + 0.10875313, + 0.01231344, + 0.099401809, + 0.10349127, + 0.10409211, + -0.1158703, + -0.063847996, + 0.034497358, + -0.059154913, + -0.041571155, + -0.050416511, + -0.097928375, + -0.046242323, + 0.014011147, + 0.16712967, + -0.03722994, + 0.058948252, + 0.048655991]], + 'cos_sim': [0.088690311, + -0.1217925, + 0.16013087, + 0.24840398, + -0.097933128, + 0.051657435, + -0.21578307, + -0.065604873, + -0.14758946], + 'dot_prods': [1.0766923, + -1.5167346, + 1.9608122, + 3.1155715, + -1.2052668, + 0.64061511, + -2.6917832, + -0.82298517, + -1.8410567], + 'score': 0.33333333}} + testing_utils.assert_deep_almost_equal( + self, expected, _unit_scale_cav(result[0]) + ) def test_tcav_ignores_labels(self): random.seed(0) # Sets seed since create_comparison_splits() uses random. @@ -192,167 +314,159 @@ def test_tcav_ignores_labels(self): ) self.assertLen(result, 1) - expected = { - 'p_val': 0.13311, - 'random_mean': 0.56667, - 'result': { - 'score': 0.33333, - 'cos_sim': [ - 0.088691, - -0.12179, - 0.16013, - 0.24840, - -0.09793, - 0.05166, - -0.21578, - -0.06560, - -0.14759, - ], - 'dot_prods': [ - 189.085096, - -266.36317, - 344.350498, - 547.144949, - -211.663965, - 112.502439, - -472.72066, - -144.529598, - -323.31888, - ], - 'accuracy': 0.66667, - 'cav': [[ - 0.91326976, - -13.67988429, - -1.33965892, - 12.23702997, - 4.83590182, - -20.6276055, - 19.07518597, - 4.9641991, - -5.52570226, - 10.37834134, - 0.53685622, - -5.21721739, - 12.00234068, - -26.36114057, - 6.20088358, - 4.76729567, - 26.4359945, - 16.30961659, - 14.70620605, - -13.47771528, - 5.29218365, - 0.87290488, - 10.19441762, - 2.96687215, - -0.70745918, - -12.201803, - -10.65010904, - -5.90814342, - -25.17510006, - -5.90629019, - 26.53638293, - 27.44054429, - -11.75430314, - -3.72779609, - 10.28197421, - -11.58444132, - -18.09351946, - -22.09520524, - -30.04023056, - -16.7551763, - -12.34913637, - 30.27897114, - -13.98790656, - 0.65758253, - 3.75770261, - 4.81132118, - -11.18005293, - 5.85445303, - 18.88336475, - -7.34885733, - -35.76094848, - -3.39495953, - 5.52774132, - -8.38126488, - 10.07413613, - -1.96825956, - 17.97041089, - -2.47085774, - 19.5700424, - -49.46295186, - 0.12541183, - -6.95592842, - -0.33562953, - 12.9269603, - -13.6288284, - -9.51211543, - 21.22778867, - -2.81344371, - -9.66434107, - -15.41551695, - -30.15406401, - 14.65903841, - 2.51729041, - 17.70711379, - 13.21615045, - 12.55899318, - 11.45749114, - -25.67659992, - 13.00876054, - -1.52499005, - 27.45026658, - -4.36110401, - 10.01664277, - -11.24470769, - -6.79411522, - -1.67106503, - -2.59389537, - 11.72310319, - -0.84126818, - 2.03886137, - -11.25047383, - -8.60679631, - -23.1676229, - 22.83532544, - -25.2657054, - -15.49795527, - -1.62890474, - -15.49504251, - 27.26973702, - -8.00652979, - 6.87541861, - -3.61878753, - 0.82889822, - -2.88891667, - 6.13730358, - 12.55884424, - -2.24121286, - 45.90285087, - 34.43108722, - -13.32567113, - 19.0988537, - 2.16242269, - 17.45654791, - 18.17472208, - 18.28023357, - -20.34869744, - -11.21275755, - 6.0583063, - -10.38857432, - -7.30056744, - -8.85395997, - -17.19779724, - -8.12087822, - 2.46058364, - 29.35061395, - -6.53820405, - 10.3522653, - 8.54478485, - ]], - }, - } - testing_utils.assert_deep_almost_equal(self, expected, result[0]) + expected = {'p_val': 0.10634367, + 'random_mean': 0.54444444, + 'result': {'accuracy': 0.66666667, + 'cav': [[0.0052003721, + -0.07789658, + -0.0076281652, + 0.069680482, + 0.02753685, + -0.11745859, + 0.10861865, + 0.028267169, + -0.031464793, + 0.059096765, + 0.0030570414, + -0.029708052, + 0.06834428, + -0.15010664, + 0.03530933, + 0.027146045, + 0.15053286, + 0.092870869, + 0.083740711, + -0.076745257, + 0.030134914, + 0.0049704607, + 0.058049336, + 0.016893858, + -0.0040285308, + -0.06947986, + -0.060644325, + -0.03364237, + -0.14335285, + -0.033631887, + 0.15110439, + 0.156253, + -0.066931739, + -0.02122692, + 0.058548115, + -0.06596458, + -0.10302878, + -0.12581527, + -0.17105629, + -0.095408045, + -0.070318855, + 0.17241567, + -0.079650335, + 0.0037445326, + 0.021397213, + 0.027396964, + -0.063661948, + 0.033336654, + 0.10752612, + -0.041846175, + -0.20363152, + -0.019331753, + 0.031476293, + -0.047724877, + 0.057364281, + -0.011207522, + 0.10232788, + -0.014069673, + 0.11143632, + -0.28165388, + 0.00071425573, + -0.039608657, + -0.0019110381, + 0.073609203, + -0.077605829, + -0.054164339, + 0.12087602, + -0.016020484, + -0.055031046, + -0.087779574, + -0.17170441, + 0.083472103, + 0.014334007, + 0.1008285, + 0.075256057, + 0.071513988, + 0.06524168, + -0.14620869, + 0.074075021, + -0.008683526, + 0.15630831, + -0.024833282, + 0.05703707, + -0.064029969, + -0.038687266, + -0.0095154466, + -0.014770232, + 0.066754095, + -0.0047905217, + 0.011609488, + -0.064062886, + -0.049009111, + -0.13192189, + 0.13002953, + -0.14386889, + -0.088248827, + -0.0092752138, + -0.088232592, + 0.1552804, + -0.045591053, + 0.039150309, + -0.020606294, + 0.0047199279, + -0.016450174, + 0.03494728, + 0.071513101, + -0.012762167, + 0.26138169, + 0.19605877, + -0.075879507, + 0.10875313, + 0.01231344, + 0.099401809, + 0.10349127, + 0.10409211, + -0.1158703, + -0.063847996, + 0.034497358, + -0.059154913, + -0.041571155, + -0.050416511, + -0.097928375, + -0.046242323, + 0.014011147, + 0.16712967, + -0.03722994, + 0.058948252, + 0.048655991]], + 'cos_sim': [0.088690311, + -0.1217925, + 0.16013087, + 0.24840398, + -0.097933128, + 0.051657435, + -0.21578307, + -0.065604873, + -0.14758946], + 'dot_prods': [1.0766923, + -1.5167346, + 1.9608122, + 3.1155715, + -1.2052668, + 0.64061511, + -2.6917832, + -0.82298517, + -1.8410567], + 'score': 0.33333333}} + testing_utils.assert_deep_almost_equal( + self, expected, _unit_scale_cav(result[0]) + ) def test_tcav_sample_from_positive(self): random.seed(0) # Sets seed since create_comparison_splits() uses random. @@ -380,66 +494,157 @@ def test_tcav_sample_from_positive(self): ) self.assertLen(result, 1) - expected = { - 'p_val': 0.80489, - 'random_mean': 0.53333, - 'result': { - 'score': - 0.8, - 'cos_sim': [ - 0.09527, -0.20442, 0.05141, 0.14985, 0.06750, -0.28244, - -0.11022, -0.14479 - ], - 'dot_prods': [ - 152.48776, -335.64998, 82.99588, 247.80113, 109.53684, - -461.81805, -181.29095, -239.47817 - ], - 'accuracy': - 1.0, - 'cav': [[ - -9.18092006e+00, -3.25423105e+00, -1.95599351e+01, - 7.90544869e+00, 2.71945760e-01, 9.21655507e+00, -7.10360094e-01, - 1.06792122e+01, -4.22297728e+00, -5.88364071e+00, - -5.90468259e+00, 3.02462186e-01, 2.73947077e+01, - -4.21289488e+00, 4.10535370e+00, 2.90032257e+00, - -1.15293947e+01, -6.05013625e+00, -9.63492785e+00, - -1.24673936e+01, -7.16846202e+00, 1.65560561e+01, - 3.91797282e+00, 9.43190766e+00, 1.13459987e+01, -3.52568932e+00, - 9.13124442e-01, 9.91573080e-01, -3.25439546e+00, - -9.93854182e-01, -1.07992122e+01, 3.11628229e+00, - 6.59499537e+00, -1.36178363e+01, 1.29858952e+01, 1.03866086e+01, - 1.59031300e+01, -5.56392889e+00, -3.77058407e+00, - -9.32599755e+00, -9.47891226e-01, 2.05803580e+01, - -1.37954357e+01, -7.73649342e+00, -9.07151538e+00, - 9.29099926e-01, -1.08088474e+01, -7.80737270e-01, - 2.40018316e+00, 1.09211064e+01, -3.27414948e+01, - -2.34077057e+01, 1.17713587e+01, 1.81635854e+00, 2.31977099e+01, - 1.61744777e+00, 1.17686967e+00, -3.88545806e+00, 2.46277754e+00, - -4.51580108e+01, 9.32061903e+00, -1.73902286e+01, - -3.66235470e+00, 3.26925824e+01, 3.21295843e+00, - -1.79316338e+01, 2.01011182e+01, 3.49235823e+00, 1.52195300e+01, - -5.04835360e+00, 1.45131574e+01, 1.59716750e+01, - -2.96747872e+00, 5.25282201e+00, -1.98576797e+01, - -1.46852052e+00, 1.58219867e+00, -5.74221070e-01, - 1.16208072e+01, 5.11250274e-01, 5.52448443e+00, 1.11949046e+01, - -6.93072443e-02, -7.17318663e+00, -8.48821209e+00, - -8.71016927e+00, -1.05849366e+01, 8.71665351e+00, - 1.24408289e+01, 4.21549502e+00, -2.07208098e+01, 1.56304334e+01, - -4.81732341e+00, 3.20792625e+01, -1.20141017e+01, - -1.73165977e+01, -6.54369967e+00, 3.16826359e+00, - -1.19574116e+01, -5.90525906e+00, 1.08200571e-01, - -9.32849433e+00, -1.29775955e+00, -1.20217051e+01, - 1.30472736e+01, 2.25283409e+01, 7.04978975e-02, 5.09369848e-01, - 1.27449879e+01, -2.16621823e+00, -8.19007901e+00, - -4.16839353e+00, 9.41498786e+00, 6.79635008e+00, 4.88519035e+00, - 2.63718274e+00, -5.28336147e-03, 1.26470921e+00, - -3.35551546e-01, 8.26617183e+00, 2.58916933e+00, - -8.54040920e+00, -1.18897963e+00, -5.09214401e+00, - 7.85278007e+00, -1.26074616e+01, -3.05422845e+00, 3.79508590e+00 - ]] - } - } - testing_utils.assert_deep_almost_equal(self, expected, result[0]) + expected = {'p_val': 0.80488552, + 'random_mean': 0.53333333, + 'result': {'accuracy': 1.0, + 'cav': [[-0.069632366, + -0.02468178, + -0.14835154, + 0.059958443, + 0.0020626273, + 0.069902726, + -0.005387655, + 0.08099582, + -0.032029033, + -0.044624209, + -0.044783667, + 0.0022943309, + 0.20777406, + -0.031952508, + 0.03113699, + 0.021997202, + -0.087444589, + -0.045886867, + -0.073075578, + -0.094558604, + -0.054369066, + 0.12556899, + 0.029715629, + 0.071535997, + 0.086053036, + -0.026740203, + 0.0069254497, + 0.0075206063, + -0.024682965, + -0.0075380765, + -0.081906185, + 0.023635468, + 0.050019659, + -0.10328389, + 0.098491028, + 0.078777097, + 0.12061685, + -0.042199314, + -0.028597893, + -0.070732988, + -0.0071890703, + 0.1560909, + -0.10463077, + -0.058677107, + -0.068802565, + 0.0070469067, + -0.081979401, + -0.0059214379, + 0.018204013, + 0.082830764, + -0.24832688, + -0.17753528, + 0.089279488, + 0.013775984, + 0.17594205, + 0.012267457, + 0.0089258831, + -0.029469196, + 0.018678911, + -0.34249923, + 0.070692293, + -0.13189565, + -0.027777011, + 0.24795561, + 0.024368627, + -0.13600187, + 0.15245624, + 0.026487686, + 0.11543197, + -0.03828894, + 0.11007448, + 0.12113642, + -0.022506757, + 0.039840035, + -0.15060978, + -0.011137922, + 0.01200007, + -0.0043553589, + 0.088137604, + 0.0038776961, + 0.041900229, + 0.084907167, + -0.00052548264, + -0.054404709, + -0.064378418, + -0.06606175, + -0.080281213, + 0.066111065, + 0.094356984, + 0.031972256, + -0.15715624, + 0.11854848, + -0.036536817, + 0.24330392, + -0.091120563, + -0.13133693, + -0.049630404, + 0.024029532, + -0.090690479, + -0.044788308, + 0.00082044973, + -0.070751615, + -0.0098426994, + -0.091178298, + 0.098956443, + 0.17086539, + 0.00053492188, + 0.0038632636, + 0.096663781, + -0.01642965, + -0.062117498, + -0.031615153, + 0.071407616, + 0.051546711, + 0.037051652, + 0.020001721, + -3.9979805e-05, + 0.0095921196, + -0.0025448876, + 0.062694438, + 0.019637428, + -0.064774334, + -0.0090176277, + -0.03862134, + 0.059559155, + -0.095620848, + -0.02316449, + 0.028783681]], + 'cos_sim': [0.095267177, + -0.20441942, + 0.051406827, + 0.14984745, + 0.067504279, + -0.2824432, + -0.11022436, + -0.1447887], + 'dot_prods': [1.1565349, + -2.5457232, + 0.62947971, + 1.8794401, + 0.83077776, + -3.5026398, + -1.3749924, + -1.8163123], + 'score': 0.8}} + testing_utils.assert_deep_almost_equal( + self, expected, _unit_scale_cav(result[0]) + ) # TODO(b/254110131): Paramaterize this test. the values returned by the # current implementaiton change when naively paramterized because of the @@ -475,67 +680,192 @@ def test_relative_tcav(self): ) self.assertLen(result, 1) - expected = { - 'result': { - 'score': 1.0, - 'cos_sim': [ - 0.9999999581246426, 0.049332143689572144, 0.8987945047547466, - -0.41858423757857954, 0.6908297036543664, -0.5167857909664919, - 0.8423017503220364, -0.005793079244916016, 0.8334491603894322, - -0.4054645113448612, 0.7616102123736647, -0.4578596155267783, - 0.8366905563807711, -0.27390786544756535, 0.7325538474066896, - 0.5190287630768531, 0.8145227936096425, 0.02005592868363552, - -0.1143256029298114, -0.1221480700842533, 0.6852995739227957, - 0.3984620730733816, 0.5211149530112407, 0.5909723902471223 - ], - 'dot_prods': [ - 1385.1480610241554, 69.95638452724207, 1239.4947646060161, - -595.253135700978, 971.5880156862692, -725.0749813217176, - 1182.8641913758102, -8.149647641120662, 1146.5803071544124, - -576.4043054391316, 1038.3510704649307, -648.097269442522, - 1154.4720122394317, -378.32103870822493, 1024.066390571124, - 738.6959135414066, 1139.7963358416857, 28.691395032352318, - -167.37808507284706, -176.4474746971391, 959.5159619261449, - 562.8772536987927, 716.7270332848395, 840.7031847912738 - ], - 'accuracy': 0.5, - 'cav': [[ - 2.03451865, -13.17604027, -8.97477873, 14.2175882, 1.75463727, - -0.46571316, 15.63798371, 9.93142109, -16.8851253, 9.11427972, - -11.06386059, 7.19105954, -17.18586585, -4.17196936, - 29.00569093, 4.33784787, 12.33925142, -15.43126877, - -24.48932802, 9.00534947, -2.19966647, -0.43898864, 10.17245833, - -1.75776027, 5.32707314, -10.89541682, -2.57798881, 4.57338285, - -1.9818453, 8.84294916, -5.4750157, 4.17024913, -19.01907373, - -1.23878054, 0.78451806, -11.60109751, 6.45912611, 1.2580972, - -11.36715828, 3.96784638, 5.01414527, -8.80780738, -17.95446165, - 6.28686056, 11.37040026, -9.99612493, -5.1080606, 9.22049438, - -6.30251447, 3.63926025, 13.14918974, 14.03763789, 5.7383655, - -8.20873775, 1.60465614, -17.91190497, -13.33602912, - -4.49996902, 28.75901011, -3.47341192, -3.82988803, 7.55923052, - -8.09407339, -16.11918298, -1.71507548, -1.36150357, - -5.86025439, 15.68441846, -3.51353798, 2.43790252, -2.19348517, - 5.0211597, 2.37646679, 13.41654008, 0.55570348, -6.05668506, - -8.42704919, -5.02851713, 14.84420539, -7.34185174, 4.03091355, - -5.09645307, 8.2008516, 1.08925377, -0.58413633, 9.92903691, - 2.14726806, 13.30179969, 13.57605882, -2.90356308, -5.45706524, - 10.79901632, -7.80617599, 17.51044709, -7.5189762, -16.52918401, - -0.52797854, -2.44961364, 10.13560444, -13.19007029, - -0.93666233, 0.54694589, 1.92814453, -4.7547722, -1.39708064, - 6.24580764, -1.41235513, -2.47243578, 10.31137604, -11.16312467, - -10.42280945, -13.08010871, 17.36711388, 5.96436519, - -11.97783485, 18.54090287, 7.91701915, -17.52516079, - -13.20368127, -26.3674698, -2.45510052, 5.56414756, 10.80934702, - -8.21923153, -11.8954632, -4.44263104, -3.42691711, -4.79007185 - ]] - }, - 'p_val': None, - 'random_mean': 0.9285714285714286, - 'split_size': 3, - 'num_runs': 1 - } + expected = {'num_runs': 1, + 'p_val': None, + 'random_mean': 0.92857143, + 'result': {'accuracy': 0.5, + 'cav': [[0.017518602, + -0.11345463, + -0.077278852, + 0.12242302, + 0.015108586, + -0.0040100901, + 0.13465355, + 0.085516162, + -0.1453923, + 0.078480072, + -0.095267333, + 0.061919861, + -0.14798188, + -0.035923496, + 0.24975854, + 0.037351765, + 0.10624928, + -0.13287364, + -0.21086958, + 0.077542178, + -0.018940547, + -0.0037800041, + 0.08759176, + -0.015135531, + 0.045869663, + -0.093816862, + -0.02219823, + 0.039379861, + -0.017064976, + 0.076143779, + -0.047143549, + 0.035908673, + -0.16376705, + -0.010666749, + 0.0067552379, + -0.099893279, + 0.055617444, + 0.010833075, + -0.097878829, + 0.034165855, + 0.04317518, + -0.075841136, + -0.15459999, + 0.054134097, + 0.097906828, + -0.086073369, + -0.043983851, + 0.079394624, + -0.054268908, + 0.031336442, + 0.11322323, + 0.12087354, + 0.049411193, + -0.070682786, + 0.013817173, + -0.1542336, + -0.11483222, + -0.038747754, + 0.24763446, + -0.029908417, + -0.032977965, + 0.065090097, + -0.069695383, + -0.13879703, + -0.01476801, + -0.011723483, + -0.050460737, + 0.13505341, + -0.030253923, + 0.020991983, + -0.018887458, + 0.043235563, + 0.020462979, + 0.1155254, + 0.0047849515, + -0.052152175, + -0.07256265, + -0.043298945, + 0.12781858, + -0.063218273, + 0.034708884, + -0.04388389, + 0.070614874, + 0.0093792398, + -0.0050298064, + 0.085495733, + 0.018489404, + 0.11453741, + 0.11689901, + -0.025001641, + -0.046989001, + 0.092986852, + -0.067216344, + 0.15077671, + -0.064743452, + -0.14232744, + -0.0045462926, + -0.021092845, + 0.087274335, + -0.11357538, + -0.0080653224, + 0.0047096042, + 0.016602641, + -0.040941779, + -0.012029788, + 0.053780589, + -0.012161309, + -0.021289336, + 0.088787906, + -0.096121974, + -0.089747414, + -0.11262854, + 0.14954251, + 0.051357184, + -0.1031372, + 0.15964969, + 0.06817089, + -0.15090343, + -0.1136926, + -0.22704156, + -0.021140095, + 0.047911059, + 0.093075737, + -0.07077305, + -0.10242786, + -0.038254045, + -0.029508093, + -0.041245751]], + 'cos_sim': [1.0000001, + 0.049331982, + 0.89879453, + -0.4185842, + 0.69082969, + -0.51678574, + 0.84230173, + -0.0057931328, + 0.83344907, + -0.4054645, + 0.76161021, + -0.45785958, + 0.83669049, + -0.27390775, + 0.7325539, + 0.51902878, + 0.8145228, + 0.020055901, + -0.11432568, + -0.12214805, + 0.68529963, + 0.398462, + 0.52111501, + 0.59097242], + 'dot_prods': [11.927058, + 0.60236967, + 10.672885, + -5.1255302, + 8.3660278, + -6.2433834, + 10.185257, + -0.07017459, + 9.8728285, + -4.9632292, + 8.9409008, + -5.5805531, + 9.9407806, + -3.2575984, + 8.8179016, + 6.3606687, + 9.814414, + 0.24705191, + -1.4412391, + -1.5193313, + 8.2620792, + 4.8467512, + 6.1715031, + 7.2390203], + 'score': 1.0}, + 'split_size': 3} - testing_utils.assert_deep_almost_equal(self, expected, result[0], places=3) + testing_utils.assert_deep_almost_equal( + self, expected, _unit_scale_cav(result[0]), places=3 + ) # This example has enough inputs for two runs of size 3. config = { @@ -567,67 +897,192 @@ def test_relative_tcav(self): config=config, ) self.assertLen(result, 1) - expected = { - 'result': { - 'score': 0.0, - 'cos_sim': [ - 0.2731987606830683, 0.427838045403812, 0.3166440584420665, - -0.1358964965831398, 0.5616614702946262, -0.16511808390168164, - -0.05103355252438478, -0.16945565920473257, 0.28148962348967155, - -0.18169036476392003, 0.33244873698665106, -0.13316476546155087, - 0.15226772288202886, -0.05534469666649352, 0.2886150002073456, - 0.33888135113008555, 0.12875301375254147, 0.046908665182593096, - -0.052445114502024985, 0.088858405172313, 0.219517174438115, - 0.35833013079793435, 0.2291162415605806, 0.3635686086637199 - ], - 'dot_prods': [ - 452.17220644153525, 724.9460578876271, 521.776546745851, - -230.9170522777958, 943.8754747127095, -276.8190148523963, - -85.63511897570154, -284.8487792023684, 462.71830216201926, - -308.62790255581496, 541.5830529968077, -225.2299308998058, - 251.04716264718752, -91.33998249705493, 482.0991668852444, - 576.3029773313335, 215.28329927312336, 80.18458502795752, - -91.74640483442752, 153.37559992294862, 367.2562273288043, - 604.8378479001944, 376.53473821563625, 618.003311205616 - ], - 'accuracy': 0.5, - 'cav': [[ - 7.55329281, -2.62460237, -9.0024543, 0.59818666, -14.58093322, - -15.42218557, -5.5354822, 12.364417, 21.63543734, -12.10592188, - -28.13137179, 14.31976269, -31.59852801, -9.5688074, - 12.95670217, -22.25037747, 16.35172203, 1.05301289, -31.3269959, - -2.59048782, 19.77513966, 10.23254963, 4.18399631, 4.87700827, - -16.33114267, 2.79754082, 10.54735581, 11.78229815, - -12.25723096, 4.70124727, -11.24257086, -5.43199386, - -7.69344958, 14.56370698, -5.75688441, -5.26137455, 23.44984724, - -6.31077994, 17.08448569, 9.05777822, 1.6895136, -17.90790926, - -14.34180658, 10.53637, 20.52315338, 13.25800162, 10.93816225, - -3.75636305, -7.45849454, -3.97231903, -2.56363646, 7.11917039, - 8.16617295, -7.89912694, 12.99981985, -21.34898036, - -17.63337856, -2.63946913, -4.40439706, 6.52652447, 3.80778434, - -9.19669592, -6.7804855, -1.69738751, 14.07850725, -19.72428768, - -4.39987548, 0.52300402, -1.32919268, -1.17373043, 25.24466652, - 1.10433416, 2.28887964, 2.73226767, 11.38988286, -15.36477422, - 1.72328432, -10.98681852, 4.09940512, 1.35212746, -11.08208377, - 15.02775212, 24.25939889, 2.5819625, -4.82504369, -2.8454034, - 24.2562686, -4.11429068, -1.74594319, -19.37503033, 2.29546159, - 6.19942707, 3.62587758, 17.79947087, 1.53031934, 2.38564874, - -16.59565939, -5.20520209, -4.84197479, -9.78838634, - -12.3678921, 16.95331818, 7.66526629, 20.39393112, 2.87282507, - -17.29492882, -7.67759675, 2.78484882, 7.79825779, -5.36932978, - 12.74328979, -5.22601154, 3.36537017, -16.2269691, 5.08010908, - 8.33516097, 6.93697789, -30.79130942, 8.65766065, 14.05914846, - 1.90387642, 7.06612428, 13.16682386, -7.97255194, -15.33183742, - 2.61200359, -2.05912859, 10.68784471 - ]] - }, - 'p_val': 0.42264973081037427, - 'random_mean': 0.0, - 'split_size': 3, - 'num_runs': 2 - } + expected = {'num_runs': 2, + 'p_val': 0.42264973, + 'random_mean': 0.0, + 'result': {'accuracy': 0.5, + 'cav': [[0.054430809, + -0.018913541, + -0.06487374, + 0.0043106256, + -0.10507367, + -0.11113577, + -0.039890099, + 0.089100905, + 0.15591004, + -0.087237962, + -0.20272122, + 0.10319157, + -0.22770646, + -0.068955034, + 0.093368895, + -0.16034138, + 0.1178343, + 0.0075883032, + -0.22574934, + -0.018667597, + 0.1425043, + 0.073738001, + 0.030150836, + 0.035144843, + -0.11768594, + 0.020159759, + 0.07600677, + 0.084905908, + -0.08832854, + 0.033878293, + -0.081016645, + -0.039144158, + -0.055440754, + 0.10494952, + -0.041485433, + -0.037914623, + 0.16898501, + -0.04547701, + 0.12311478, + 0.06527245, + 0.012174988, + -0.1290486, + -0.10335032, + 0.075927481, + 0.14789458, + 0.095540419, + 0.078822859, + -0.027069161, + -0.053747706, + -0.028625416, + -0.018474095, + 0.051302161, + 0.058847334, + -0.056922939, + 0.093679711, + -0.15384541, + -0.12707013, + -0.019020651, + -0.031739205, + 0.047031708, + 0.02743979, + -0.066273555, + -0.048861716, + -0.012231816, + 0.10145299, + -0.14213774, + -0.031706437, + 0.0037689353, + -0.0095784897, + -0.0084581785, + 0.1819189, + 0.0079581132, + 0.016494168, + 0.019689338, + 0.082078144, + -0.11072212, + 0.012418305, + -0.07917355, + 0.029541299, + 0.0097436868, + -0.079859987, + 0.10829336, + 0.17481889, + 0.018606279, + -0.03477037, + -0.020504609, + 0.17479621, + -0.029648563, + -0.012581717, + -0.13962108, + 0.016541589, + 0.044674456, + 0.026129019, + 0.12826715, + 0.011027789, + 0.017191606, + -0.1195922, + -0.037509825, + -0.034892451, + -0.070537373, + -0.08912605, + 0.12216952, + 0.055237662, + 0.14696342, + 0.020702258, + -0.12463128, + -0.055326533, + 0.020068195, + 0.056196019, + -0.038692404, + 0.091831148, + -0.037659943, + 0.024251722, + -0.11693534, + 0.036608536, + 0.060065098, + 0.049989466, + -0.22188933, + 0.062389087, + 0.10131354, + 0.013719752, + 0.050920133, + 0.094882987, + -0.057452019, + -0.11048462, + 0.018822581, + -0.014838634, + 0.077019162]], + 'cos_sim': [0.27319872, + 0.427838, + 0.31664395, + -0.13589641, + 0.56166136, + -0.16511787, + -0.051033769, + -0.16945563, + 0.28148949, + -0.18169014, + 0.33244866, + -0.13316464, + 0.15226755, + -0.055344634, + 0.28861475, + 0.33888116, + 0.12875293, + 0.046908617, + -0.052445047, + 0.088858448, + 0.21951714, + 0.35833022, + 0.22911605, + 0.36356851], + 'dot_prods': [3.2584567, + 5.2241292, + 3.7600415, + -1.6640407, + 6.801784, + -1.9948194, + -0.61710912, + -2.0526855, + 3.3344538, + -2.2240415, + 3.902771, + -1.6230575, + 1.8091018, + -0.65821648, + 3.4741149, + 4.1529698, + 1.5513802, + 0.57782811, + -0.66114497, + 1.1052606, + 2.646533, + 4.3586025, + 2.7133942, + 4.4534731], + 'score': 0.0}, + 'split_size': 3} - testing_utils.assert_deep_almost_equal(self, expected, result[0], places=3) + testing_utils.assert_deep_almost_equal( + self, expected, _unit_scale_cav(result[0]), places=3 + ) # This example has enough examples for three runs of size 3 and two runs of # size 5, and returns results with p-value < 0.05. @@ -668,79 +1123,193 @@ def test_relative_tcav(self): config=config, ) self.assertLen(result, 1) - expected = [{ - 'result': { - 'score': 0.42857142857142855, - 'cos_sim': [ - -0.1107393877916321, -0.0993967046974328, -0.2214985917242054, - 0.08132588965575606, -0.3590211572508748, 0.18708109817461333, - 0.000724498781128839, 0.09700473783330398, -0.25015742815240055, - 0.16108236033785076, -0.10283274286140846, 0.0972663321478731, - -0.05924679176256152, -0.048499696342091746, - -0.4357117016074766, -0.593245752003111, -0.3645147796989344, - -0.5507605083253673, -0.27914997949782694, -0.30908550968594417, - -0.5584676299422896, -0.16983339994284577, -0.42587740852240746, - -0.37482298817032594 - ], - 'dot_prods': [ - -261.4389298435066, -240.23776409902007, -520.6275907607769, - 197.11495117497446, -860.6035066083074, 447.3775519523981, - 1.7341104803878409, 232.59170976304426, -586.5576327736542, - 390.2961568516803, -238.95427152619726, 234.6617547723058, - -139.3334215524385, -114.17392512371171, -1038.149036709951, - -1439.0663895591745, -869.3828698612926, -1342.899780229334, - -696.569760699206, -760.9907977738051, -1332.7284530349625, - -408.90435403478875, -998.3360993150825, -908.8111404537224 - ], - 'accuracy': 0.75, - 'cav': [[ - -2.01128835e+01, -2.53418963e+01, 3.02074307e+01, - 3.78726319e+00, -6.30370696e+00, 1.74700275e+00, 2.26769890e+01, - -1.06487415e+01, -1.76710295e+00, 1.53911930e+01, - 6.54822833e+00, 2.02141240e+00, 1.16919236e+00, -4.59410912e-01, - 4.54187735e+00, 1.98659207e+01, 1.01712167e-01, 1.18509678e+01, - 3.11218189e+01, 1.60491967e+01, -2.84655620e+01, - -1.08942472e+01, -1.64852881e+01, 4.79670684e+00, - 3.08027345e+01, -2.64932823e+01, 1.14056950e+00, - -1.45201048e+01, -8.41149151e+00, -1.26598740e+01, - -5.87020201e+00, -2.87027220e+00, 2.13724773e+01, - 3.16488669e+01, -1.38671518e+01, 2.25443206e+01, - -1.93658542e+01, 1.47272427e+01, 1.42369662e+01, - -7.98238134e+00, -2.34073887e+01, 3.63425744e+00, - 3.20426643e+00, -1.57109214e+01, -1.87085664e+01, - 2.04525909e-02, 1.75499187e+01, -1.89476242e+00, 1.22602038e+01, - -8.30212717e-01, 6.89138456e+00, -4.76848738e+01, - -1.02832284e+01, 1.25879647e+01, -2.34436812e+01, - 3.46227988e+01, 3.98466501e+00, -4.18899781e+01, 2.31382829e+01, - 9.44899112e+00, -5.55529692e+00, 2.99535098e+01, 1.91656336e+01, - 1.90906964e+01, -1.97222943e+01, -8.81455291e+00, - -4.31148519e-01, 1.35055378e+01, -2.93433908e+01, - 1.41053039e+01, 7.31250277e+00, 1.84754657e+01, -3.59438637e+00, - -1.94830172e+01, 1.82690513e+00, 1.11006786e+01, - -1.98780467e+01, 9.70483509e+00, 3.10464830e+01, - -1.58164540e+00, 3.92983591e+01, -1.94148650e+01, - -1.78218017e+01, 1.68395976e+01, 1.88509800e+01, - -4.67428605e+00, -4.11605219e+01, 4.20027328e+00, - -1.78356055e+01, 2.93033407e+01, 3.67970865e+00, 4.43803533e+00, - -5.86478123e+00, -1.89515809e+01, -4.15883069e+00, - -2.53725359e+01, -2.06732149e+01, 3.51060797e+00, - 7.32087133e+00, 7.98146613e-01, 3.38068707e+00, -2.48928693e+01, - 1.99907796e+00, 2.12458621e+00, -1.33852676e+01, 1.42904487e+01, - 1.33286404e-02, -3.82017252e+00, -1.04635257e+01, - 1.00280360e+01, 8.93273804e+00, 5.26520675e+00, -1.28270015e+01, - 3.54992057e+00, 2.32608393e+00, -1.84115520e+01, - -1.74357882e+00, 2.53797871e+01, -1.75258907e+01, - -4.22843710e+01, -1.49635329e+01, -3.77852872e+00, - -1.25404291e+01, 1.17741455e+01, 3.35462077e+01, 1.65145075e+01, - 4.27775903e-01, 8.26403983e+00 - ]] - }, - 'p_val': 0.04400624968940752, - 'random_mean': 0.9642857142857143, - 'split_size': 5, - 'num_runs': 2 - }] - testing_utils.assert_deep_almost_equal(self, expected, result, places=3) + expected = [ + [{'num_runs': 2, + 'p_val': 0.049345849, + 'random_mean': 0.60714286, + 'result': {'accuracy': 0.5, + 'cav': [[-0.055570453, + -0.096331999, + 0.012094902, + -0.08287356, + -0.10415225, + 0.21779795, + -0.072624646, + 0.15025093, + -0.073818751, + -0.12635337, + -0.046921663, + -0.052306481, + 0.0055506672, + -0.048594911, + 0.0056963833, + 0.088665105, + -0.018665945, + 0.047776558, + -0.10221633, + -0.024074927, + -0.048519857, + -0.010542592, + -0.19978398, + -0.025263591, + -0.16287366, + -0.020170791, + 0.045984548, + 0.14284854, + -0.054188672, + -0.09720853, + -0.050516028, + 0.074546874, + -0.03794039, + 0.06347952, + -0.010381403, + -0.06778311, + -0.050439533, + -0.066106305, + 0.032643512, + -0.081726588, + 0.056773905, + -0.091130488, + -0.02799716, + 0.10021219, + -0.090247974, + 0.079484731, + 0.091580577, + -0.043116856, + -0.043661203, + 0.024227265, + 0.11165099, + 0.12508449, + 0.066936687, + 0.034734819, + -0.040805452, + -0.037688263, + 0.083753683, + 0.027496248, + -0.098193742, + 0.11788043, + -0.066893943, + -0.058402538, + -0.074292183, + -0.031210374, + 0.021727433, + -0.22137223, + 0.045704812, + 0.040843945, + 0.13677552, + -0.012557695, + 0.17133912, + -0.10822944, + 0.060261216, + 0.085893132, + 0.018889453, + 0.030267421, + 0.0460894, + -0.10867414, + 0.02763474, + -0.092986099, + 0.028839169, + 0.18626508, + 0.036116894, + 0.0034696588, + -0.063530527, + 0.02104692, + 0.045171186, + -0.02538478, + 0.097192526, + -0.17315385, + -0.060331281, + 0.019191742, + 0.21491741, + -0.10249273, + 0.059563346, + -0.17183439, + 0.050312571, + 0.2300659, + 0.037521385, + -0.082064256, + -0.023227613, + 0.0082738642, + 0.042664625, + -0.038796607, + 0.10451597, + -0.038144995, + 0.0511855, + -0.15643859, + -0.022560269, + 0.040286388, + 0.091795728, + -0.0045317137, + 0.13958201, + -0.0077924212, + 0.034880966, + 0.077661686, + -0.045588546, + -0.019975496, + 0.056513544, + -0.0098233251, + 0.18177171, + 0.16408995, + 0.021294123, + -0.029079312, + -0.15497734, + 0.028431851, + 0.10945288, + -0.0056610866]], + 'cos_sim': [0.065847836, + 0.38948673, + 0.045807634, + 0.27994305, + 0.12873483, + 0.15351148, + -0.046763059, + 0.16203874, + -0.067151539, + 0.16850159, + -0.42353198, + -0.11642149, + -0.32443368, + 0.043689668, + 0.17103073, + 0.26918238, + -0.044799764, + 0.17242718, + 0.1624705, + -0.047080301, + -0.0097219115, + 0.02543349, + -0.098340467, + 0.37685275], + 'dot_prods': [0.78537089, + 4.7558393, + 0.54395038, + 3.42788, + 1.5589935, + 1.8546003, + -0.56546694, + 1.9628416, + -0.79546022, + 2.0626023, + -4.9720411, + -1.418986, + -3.8546202, + 0.51960337, + 2.0587318, + 3.2988148, + -0.53980494, + 2.1239865, + 2.0481737, + -0.58560556, + -0.11720888, + 0.30936402, + -1.1646345, + 4.6161962], + 'score': 0.14285714}, + 'split_size': 5}] + ] + testing_utils.assert_deep_almost_equal( + self, expected, [_unit_scale_cav(r) for r in result], places=3 + ) if __name__ == '__main__': diff --git a/lit_nlp/components/tcav_test.py b/lit_nlp/components/tcav_test.py index 7402a5a8..1e7ad0ff 100644 --- a/lit_nlp/components/tcav_test.py +++ b/lit_nlp/components/tcav_test.py @@ -228,18 +228,28 @@ def test_compute_tcav_score(self): self.assertAlmostEqual(0.5, result) def test_get_trained_cav(self): + # CAV magnitudes depend on the SVM solver's regularization scaling, which + # varies across scikit-learn versions, so compare unit-normalized + # directions. + def _unit(cav): + return np.asarray(cav) / np.linalg.norm(cav) + # 1D inputs. x = [[1], [1], [1], [2], [1], [1], [-1], [-1], [-2], [-1], [-1]] y = [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0] cav, accuracy = self.tcav.get_trained_cav(x, y, 0.33, random_state=0) - np.testing.assert_almost_equal(np.array([[19.08396947]]), cav) + np.testing.assert_almost_equal( + _unit(np.array([[19.08396947]])), _unit(cav) + ) self.assertAlmostEqual(1.0, accuracy) # 2D inputs. x = [[-8, 1], [5, 3], [3, 6], [-2, 5], [-8, 10], [10, -5]] y = [1, 0, 0, 1, 1, 0] cav, accuracy = self.tcav.get_trained_cav(x, y, 0.33, random_state=0) - np.testing.assert_almost_equal(np.array([[-77.89678676, 9.73709834]]), cav) + np.testing.assert_almost_equal( + _unit(np.array([[-77.89678676, 9.73709834]])), _unit(cav) + ) self.assertAlmostEqual(1.0, accuracy) def test_compute_local_scores(self): diff --git a/lit_nlp/components/thresholder_int_test.py b/lit_nlp/components/thresholder_int_test.py index 9775b79c..9d8fa38f 100644 --- a/lit_nlp/components/thresholder_int_test.py +++ b/lit_nlp/components/thresholder_int_test.py @@ -14,6 +14,8 @@ # ============================================================================== """Tests for lit_nlp.components.thresholder.""" +import os + from absl.testing import absltest from absl.testing import parameterized from lit_nlp.api import dataset as lit_dataset @@ -22,12 +24,10 @@ from lit_nlp.examples.glue import models as glue_models from lit_nlp.lib import caching # for hash id fn - -from lit_nlp.lib import file_cache -BERT_TINY_PATH = file_cache.cached_path( - 'https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz', # pylint: disable=line-too-long - extract_compressed_file=True, +_TESTDATA_PATH = os.path.join( + os.path.dirname(glue_models.__file__), 'testdata' ) +BERT_TINY_PATH = os.path.join(_TESTDATA_PATH, 'sst2_tiny') _EXAMPLES = [ diff --git a/lit_nlp/examples/blank_slate_demo.py b/lit_nlp/examples/blank_slate_demo.py index c82ffde1..fa8020f6 100644 --- a/lit_nlp/examples/blank_slate_demo.py +++ b/lit_nlp/examples/blank_slate_demo.py @@ -2,10 +2,10 @@ The currently supported models and datasets are: - classification model on SST-2, with the Stanford Sentiment Treebank dataset. -- regression model on STS-B, with Semantic Textual Similarit Benchmark dataset. +- regression model on STS-B, with Semantic Textual Similarity Benchmark dataset. - classification model on MultiNLI, with the MultiNLI dataset. -- TensorFlow Keras model for penguin classification, with the Penguin tabular - dataset from TFDS. +- Scikit-learn model for penguin classification, with the Palmer Penguins + tabular dataset. To run: python -m lit_nlp.examples.blank_slate_demo --port=5432 diff --git a/lit_nlp/examples/custom_module/potato_demo.py b/lit_nlp/examples/custom_module/potato_demo.py index 6a3416ba..4125cb29 100644 --- a/lit_nlp/examples/custom_module/potato_demo.py +++ b/lit_nlp/examples/custom_module/potato_demo.py @@ -35,7 +35,7 @@ _MODEL = flags.DEFINE_string( "model", - "https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz", + "philschmid/tiny-bert-sst2-distilled", "Path to model, as in examples/glue/demo.py") # Use our custom frontend build from this directory. diff --git a/lit_nlp/examples/gcp/Dockerfile b/lit_nlp/examples/gcp/Dockerfile index 86217e1d..f7feb067 100644 --- a/lit_nlp/examples/gcp/Dockerfile +++ b/lit_nlp/examples/gcp/Dockerfile @@ -93,14 +93,14 @@ COPY ./lit_nlp/examples/gcp/model_server_gunicorn_config.py ./ FROM lit-gcp-app-server-base AS lit-gcp-app-server -RUN python -m pip install 'lit-nlp[examples-generative-ai]' +RUN python -m pip install 'lit-nlp[examples-common,examples-generative-ai]' ENTRYPOINT ["gunicorn", "--config=gunicorn_config.py"] FROM lit-gcp-model-server-base AS lit-gcp-model-server -RUN python -m pip install 'lit-nlp[examples-generative-ai]' +RUN python -m pip install 'lit-nlp[examples-common,examples-generative-ai]' ENTRYPOINT ["gunicorn", "--config=model_server_gunicorn_config.py"] diff --git a/lit_nlp/examples/gcp/README.md b/lit_nlp/examples/gcp/README.md index b27ed5e7..a43a1985 100644 --- a/lit_nlp/examples/gcp/README.md +++ b/lit_nlp/examples/gcp/README.md @@ -27,6 +27,6 @@ docker build -f ./lit_nlp/examples/gcp/Dockerfile -t lit-app:gcp-dev . ### Run GPT-2 in a Docker container ```shell -# Runs GPT-2 in Keras on PyTorch -docker run --rm -p 5432:5432 -e MODEL_CONFIG=gpt2:gpt2_base_en lit-app:gcp-dev +# Runs GPT-2 on PyTorch +docker run --rm -p 5432:5432 -e MODEL_CONFIG=gpt2:gpt2 lit-app:gcp-dev ``` diff --git a/lit_nlp/examples/gcp/model_server.py b/lit_nlp/examples/gcp/model_server.py index 4be896ae..55188f2f 100644 --- a/lit_nlp/examples/gcp/model_server.py +++ b/lit_nlp/examples/gcp/model_server.py @@ -26,7 +26,7 @@ from lit_nlp.lib import serialize from lit_nlp.lib import wsgi_app -DEFAULT_MODELS = 'gemma_1.1_2b_IT:gemma_1.1_instruct_2b_en' +DEFAULT_MODELS = 'gemma_2b_IT:google/gemma-2b-it' _LlmHTTPEndpoints = lit_gcp_constants.LlmHTTPEndpoints @@ -52,7 +52,6 @@ def _handler(wsgiapp: wsgi_app.App, request, unused_environ): f'Only 1 model configuration can be provided, got {num_configs}' ) - dl_framework = os.getenv('DL_FRAMEWORK', pd_models.DEFAULT_DL_FRAMEWORK) dl_runtime = os.getenv('DL_RUNTIME', pd_models.DEFAULT_DL_RUNTIME) precision = os.getenv('PRECISION', pd_models.DEFAULT_PRECISION) batch_size = int(os.getenv('BATCH_SIZE', pd_models.DEFAULT_BATCH_SIZE)) @@ -62,7 +61,6 @@ def _handler(wsgiapp: wsgi_app.App, request, unused_environ): models = pd_models.get_models( models_config=model_config, - dl_framework=dl_framework, dl_runtime=dl_runtime, precision=precision, batch_size=batch_size, diff --git a/lit_nlp/examples/glue/data.py b/lit_nlp/examples/glue/data.py index c66f58aa..0e3cfd7d 100644 --- a/lit_nlp/examples/glue/data.py +++ b/lit_nlp/examples/glue/data.py @@ -1,47 +1,38 @@ -"""GLUE benchmark datasets, using TFDS or from CSV. +"""GLUE benchmark datasets, from the HuggingFace datasets hub. See https://gluebenchmark.com/ and -https://www.tensorflow.org/datasets/catalog/glue +https://huggingface.co/datasets/glue -Note that this requires the TensorFlow Datasets package, but the resulting LIT -datasets just contain regular Python/NumPy data. +The resulting LIT datasets contain regular Python/NumPy data. """ from typing import Optional +from datasets import load_dataset from lit_nlp.api import dataset as lit_dataset from lit_nlp.api import types as lit_types from lit_nlp.lib import file_cache from lit_nlp.lib import utils import pandas as pd -import tensorflow_datasets as tfds -def load_tfds(*args, do_sort=True, **kw): - """Load from TFDS, with optional sorting.""" - # Materialize to NumPy arrays. - # This also ensures compatibility with TF1.x non-eager mode, which doesn't - # support direct iteration over a tf.data.Dataset. - ret = list(tfds.as_numpy(tfds.load(*args, download=True, try_gcs=True, **kw))) - if do_sort: - # Recover original order, as if you loaded from a TSV file. - ret.sort(key=lambda ex: ex['idx']) - return ret +def load_hf_glue(config_name: str, split: str, **kw): + """Load a GLUE config from HuggingFace, preserving original example order.""" + return list(load_dataset('glue', config_name, split=split, **kw)) class CoLAData(lit_dataset.Dataset): """Corpus of Linguistic Acceptability. - See - https://www.tensorflow.org/datasets/catalog/glue#gluecola_default_config. + See https://huggingface.co/datasets/glue/viewer/cola. """ LABELS = ['0', '1'] def __init__(self, split: str): self._examples = [] - for ex in load_tfds('glue/cola', split=split): + for ex in load_hf_glue('cola', split=split): self._examples.append({ - 'sentence': ex['sentence'].decode('utf-8'), + 'sentence': ex['sentence'], 'label': self.LABELS[ex['label']], }) @@ -55,11 +46,11 @@ def spec(self): class SST2Data(lit_dataset.Dataset): """Stanford Sentiment Treebank, binary version (SST-2). - See https://www.tensorflow.org/datasets/catalog/glue#gluesst2. + See https://huggingface.co/datasets/glue/viewer/sst2. """ LABELS = ['0', '1'] - TFDS_SPLITS = ['test', 'train', 'validation'] + SPLITS = ['test', 'train', 'validation'] def load_from_csv(self, path: str): path = file_cache.cached_path(path) @@ -73,15 +64,15 @@ def load_from_csv(self, path: str): df['label'] = df.label.map(str) return df.to_dict(orient='records') - def load_from_tfds(self, split: str): - if split not in self.TFDS_SPLITS: + def load_from_hf(self, split: str): + if split not in self.SPLITS: raise ValueError( - f"Unsupported split '{split}'. Allowed values: {self.TFDS_SPLITS}" + f"Unsupported split '{split}'. Allowed values: {self.SPLITS}" ) ret = [] - for ex in load_tfds('glue/sst2', split=split): + for ex in load_hf_glue('sst2', split=split): ret.append({ - 'sentence': ex['sentence'].decode('utf-8'), + 'sentence': ex['sentence'], 'label': self.LABELS[ex['label']], }) return ret @@ -92,7 +83,7 @@ def __init__( if path_or_splitname.endswith('.csv'): self._examples = self.load_from_csv(path_or_splitname)[:max_examples] else: - self._examples = self.load_from_tfds(path_or_splitname)[:max_examples] + self._examples = self.load_from_hf(path_or_splitname)[:max_examples] @classmethod def init_spec(cls) -> lit_types.Spec: @@ -115,7 +106,7 @@ def spec(self): class SST2DataForLM(SST2Data): """Stanford Sentiment Treebank, binary version (SST-2). - See https://www.tensorflow.org/datasets/catalog/glue#gluesst2. + See https://huggingface.co/datasets/glue/viewer/sst2. This data is reformatted to serve the language models. """ @@ -146,17 +137,17 @@ def init_spec(cls) -> lit_types.Spec: class MRPCData(lit_dataset.Dataset): """Microsoft Research Paraphrase Corpus. - See https://www.tensorflow.org/datasets/catalog/glue#gluemrpc. + See https://huggingface.co/datasets/glue/viewer/mrpc. """ LABELS = ['0', '1'] def __init__(self, split: str): self._examples = [] - for ex in load_tfds('glue/mrpc', split=split): + for ex in load_hf_glue('mrpc', split=split): self._examples.append({ - 'sentence1': ex['sentence1'].decode('utf-8'), - 'sentence2': ex['sentence2'].decode('utf-8'), + 'sentence1': ex['sentence1'], + 'sentence2': ex['sentence2'], 'label': self.LABELS[ex['label']], }) @@ -171,17 +162,17 @@ def spec(self): class QQPData(lit_dataset.Dataset): """Quora Question Pairs. - See https://www.tensorflow.org/datasets/catalog/glue#glueqqp. + See https://huggingface.co/datasets/glue/viewer/qqp. """ LABELS = ['0', '1'] def __init__(self, split: str): self._examples = [] - for ex in load_tfds('glue/qqp', split=split): + for ex in load_hf_glue('qqp', split=split): self._examples.append({ - 'question1': ex['question1'].decode('utf-8'), - 'question2': ex['question2'].decode('utf-8'), + 'question1': ex['question1'], + 'question2': ex['question2'], 'label': self.LABELS[ex['label']], }) @@ -198,9 +189,9 @@ class STSBData(lit_dataset.Dataset): Unlike the other GLUE tasks, this is formulated as a regression problem. - See https://www.tensorflow.org/datasets/catalog/glue#gluestsb. + See https://huggingface.co/datasets/glue/viewer/stsb. """ - TFDS_SPLITS = ['test', 'train', 'validation'] + SPLITS = ['test', 'train', 'validation'] def load_from_csv(self, path: str): path = file_cache.cached_path(path) @@ -214,16 +205,16 @@ def load_from_csv(self, path: str): df['label'] = df.label.map(float) return df.to_dict(orient='records') - def load_from_tfds(self, split: str): - if split not in self.TFDS_SPLITS: + def load_from_hf(self, split: str): + if split not in self.SPLITS: raise ValueError( - f"Unsupported split '{split}'. Allowed values: {self.TFDS_SPLITS}" + f"Unsupported split '{split}'. Allowed values: {self.SPLITS}" ) ret = [] - for ex in load_tfds('glue/stsb', split=split): + for ex in load_hf_glue('stsb', split=split): ret.append({ - 'sentence1': ex['sentence1'].decode('utf-8'), - 'sentence2': ex['sentence2'].decode('utf-8'), + 'sentence1': ex['sentence1'], + 'sentence2': ex['sentence2'], 'label': ex['label'], }) return ret @@ -234,7 +225,7 @@ def __init__( if path_or_splitname.endswith('.csv'): self._examples = self.load_from_csv(path_or_splitname)[:max_examples] else: - self._examples = self.load_from_tfds(path_or_splitname)[:max_examples] + self._examples = self.load_from_hf(path_or_splitname)[:max_examples] @classmethod def init_spec(cls) -> lit_types.Spec: @@ -258,11 +249,11 @@ def spec(self): class MNLIData(lit_dataset.Dataset): """MultiNLI dataset. - See https://www.tensorflow.org/datasets/catalog/glue#gluemnli. + See https://huggingface.co/datasets/glue/viewer/mnli. """ LABELS = ['entailment', 'neutral', 'contradiction'] - TFDS_SPLITS = [ + SPLITS = [ 'test_matched', 'test_mismatched', 'train', @@ -282,16 +273,16 @@ def load_from_csv(self, path: str): df['label'] = df.label.map(str) return df.to_dict(orient='records') - def load_from_tfds(self, split: str): - if split not in self.TFDS_SPLITS: + def load_from_hf(self, split: str): + if split not in self.SPLITS: raise ValueError( - f"Unsupported split '{split}'. Allowed values: {self.TFDS_SPLITS}" + f"Unsupported split '{split}'. Allowed values: {self.SPLITS}" ) ret = [] - for ex in load_tfds('glue/mnli', split=split): + for ex in load_hf_glue('mnli', split=split): ret.append({ - 'premise': ex['premise'].decode('utf-8'), - 'hypothesis': ex['hypothesis'].decode('utf-8'), + 'premise': ex['premise'], + 'hypothesis': ex['hypothesis'], 'label': self.LABELS[ex['label']], }) return ret @@ -302,7 +293,7 @@ def __init__( if path_or_splitname.endswith('.csv'): self._examples = self.load_from_csv(path_or_splitname)[:max_examples] else: - self._examples = self.load_from_tfds(path_or_splitname)[:max_examples] + self._examples = self.load_from_hf(path_or_splitname)[:max_examples] @classmethod def init_spec(cls) -> lit_types.Spec: @@ -326,17 +317,17 @@ def spec(self): class QNLIData(lit_dataset.Dataset): """NLI examples derived from SQuAD. - See https://www.tensorflow.org/datasets/catalog/glue#glueqnli. + See https://huggingface.co/datasets/glue/viewer/qnli. """ LABELS = ['entailment', 'not_entailment'] def __init__(self, split: str): self._examples = [] - for ex in load_tfds('glue/qnli', split=split): + for ex in load_hf_glue('qnli', split=split): self._examples.append({ - 'question': ex['question'].decode('utf-8'), - 'sentence': ex['sentence'].decode('utf-8'), + 'question': ex['question'], + 'sentence': ex['sentence'], 'label': self.LABELS[ex['label']], }) @@ -351,17 +342,17 @@ def spec(self): class RTEData(lit_dataset.Dataset): """Recognizing Textual Entailment. - See https://www.tensorflow.org/datasets/catalog/glue#gluerte. + See https://huggingface.co/datasets/glue/viewer/rte. """ LABELS = ['entailment', 'not_entailment'] def __init__(self, split: str): self._examples = [] - for ex in load_tfds('glue/rte', split=split): + for ex in load_hf_glue('rte', split=split): self._examples.append({ - 'sentence1': ex['sentence1'].decode('utf-8'), - 'sentence2': ex['sentence2'].decode('utf-8'), + 'sentence1': ex['sentence1'], + 'sentence2': ex['sentence2'], 'label': self.LABELS[ex['label']], }) @@ -376,17 +367,17 @@ def spec(self): class WNLIData(lit_dataset.Dataset): """Winograd schema challenge. - See https://www.tensorflow.org/datasets/catalog/glue#gluewnli. + See https://huggingface.co/datasets/glue/viewer/wnli. """ LABELS = ['0', '1'] def __init__(self, split: str): self._examples = [] - for ex in load_tfds('glue/wnli', split=split): + for ex in load_hf_glue('wnli', split=split): self._examples.append({ - 'sentence1': ex['sentence1'].decode('utf-8'), - 'sentence2': ex['sentence2'].decode('utf-8'), + 'sentence1': ex['sentence1'], + 'sentence2': ex['sentence2'], 'label': self.LABELS[ex['label']], }) @@ -401,17 +392,17 @@ def spec(self): class DiagnosticNLIData(lit_dataset.Dataset): """NLI diagnostic set; use to evaluate models trained on MultiNLI. - See https://www.tensorflow.org/datasets/catalog/glue#glueax. + See https://huggingface.co/datasets/glue/viewer/ax. """ LABELS = ['entailment', 'neutral', 'contradiction'] def __init__(self, split: str): self._examples = [] - for ex in load_tfds('glue/ax', split=split): + for ex in load_hf_glue('ax', split=split): self._examples.append({ - 'premise': ex['premise'].decode('utf-8'), - 'hypothesis': ex['hypothesis'].decode('utf-8'), + 'premise': ex['premise'], + 'hypothesis': ex['hypothesis'], 'label': self.LABELS[ex['label']], }) diff --git a/lit_nlp/examples/glue/demo.py b/lit_nlp/examples/glue/demo.py index 0d577032..67618d91 100644 --- a/lit_nlp/examples/glue/demo.py +++ b/lit_nlp/examples/glue/demo.py @@ -38,17 +38,16 @@ _MODELS = flags.DEFINE_list( "models", [ - "sst2-tiny:sst2:https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz", - "sst2-base:sst2:https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_base.tar.gz", - "stsb:stsb:https://storage.googleapis.com/what-if-tool-resources/lit-models/stsb_base.tar.gz", - "mnli:mnli:https://storage.googleapis.com/what-if-tool-resources/lit-models/mnli_base.tar.gz", + "sst2:sst2:textattack/bert-base-uncased-SST-2", + "stsb:stsb:Alireza1044/albert-base-v2-stsb", + "mnli:mnli:textattack/bert-base-uncased-MNLI", ], ( "List of models to load, as ::. See MODELS_BY_TASK" - " for available tasks. Path should be the output of saving a" - " transformers model, e.g. model.save_pretrained(path) and" - " tokenizer.save_pretrained(path). Remote .tar.gz files will be" - " downloaded and cached locally." + " for available tasks. Path should be a HuggingFace model name, or the" + " output of saving a transformers model, e.g." + " model.save_pretrained(path) and tokenizer.save_pretrained(path)." + " Remote .tar.gz files will be downloaded and cached locally." ), ) @@ -68,10 +67,10 @@ # Pre-specified set of small models, which will load and run much faster. QUICK_START_MODELS = ( - "sst2-tiny:sst2:https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz", - "sst2-small:sst2:https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_small.tar.gz", - "stsb-tiny:stsb:https://storage.googleapis.com/what-if-tool-resources/lit-models/stsb_tiny.tar.gz", - "mnli-small:mnli:https://storage.googleapis.com/what-if-tool-resources/lit-models/mnli_small.tar.gz", + "sst2-tiny:sst2:philschmid/tiny-bert-sst2-distilled", + "sst2-base:sst2:textattack/bert-base-uncased-SST-2", + "stsb:stsb:Alireza1044/albert-base-v2-stsb", + "mnli:mnli:textattack/bert-base-uncased-MNLI", ) diff --git a/lit_nlp/examples/glue/model_int_test.py b/lit_nlp/examples/glue/model_int_test.py index 366a9428..f3c540bb 100644 --- a/lit_nlp/examples/glue/model_int_test.py +++ b/lit_nlp/examples/glue/model_int_test.py @@ -8,11 +8,13 @@ --guitar_detach """ +import os from typing import Any from absl.testing import absltest from absl.testing import parameterized from lit_nlp.examples.glue import models as glue_models -from lit_nlp.lib import file_cache + +TESTDATA_PATH = os.path.join(os.path.dirname(__file__), "testdata") # TODO(b/254110131): Fix test flakiness. Expand to SST-2, STS-B, and MNLI @@ -21,11 +23,7 @@ class GlueModelsIntTest(parameterized.TestCase): def __init__(self, *args: Any, **kwargs: Any): super().__init__(*args, **kwargs) # Create the SST-2 Model - model_path = "https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz" # pylint: disable=line-too-long - if model_path.endswith(".tar.gz"): - model_path = file_cache.cached_path( - model_path, extract_compressed_file=True - ) + model_path = os.path.join(TESTDATA_PATH, "sst2_tiny") self.sst2_model = glue_models.SST2Model(model_path) @parameterized.named_parameters( diff --git a/lit_nlp/examples/glue/model_utils.py b/lit_nlp/examples/glue/model_utils.py index b4ddd446..6716e65c 100644 --- a/lit_nlp/examples/glue/model_utils.py +++ b/lit_nlp/examples/glue/model_utils.py @@ -1,26 +1,14 @@ """Helpers for working with transformers models.""" from typing import Optional -from absl import logging import transformers -def load_pretrained(cls, *args, **kw): - """Load a transformers model in TF2, with fallback to PyTorch weights.""" - try: - return cls.from_pretrained(*args, **kw) - except OSError as e: - logging.warning("Caught OSError loading model: %s", e) - logging.warning( - "Re-trying to convert from PyTorch checkpoint (from_pt=True)") - return cls.from_pretrained(*args, from_pt=True, **kw) - - def batch_encode_pretokenized( tokenizer: transformers.PreTrainedTokenizerBase, tokenized_inputs: list[list[str]], tokenized_pair_inputs: Optional[list[list[str]]] = None, - tensor_type="tf", + tensor_type="pt", **kw ) -> transformers.BatchEncoding: """Batch encode pre-tokenized text, without further splitting. diff --git a/lit_nlp/examples/glue/model_utils_test.py b/lit_nlp/examples/glue/model_utils_test.py index 27747a78..988cca76 100644 --- a/lit_nlp/examples/glue/model_utils_test.py +++ b/lit_nlp/examples/glue/model_utils_test.py @@ -74,7 +74,7 @@ def test_untokenized_single(self): encoded = model_utils.batch_encode_pretokenized(self.tokenizer, tokens) expected_encoded = self.tokenizer.batch_encode_plus( sentences, - return_tensors='tf', + return_tensors='pt', add_special_tokens=True, padding='longest', truncation='longest_first') @@ -94,7 +94,7 @@ def test_untokenized_pairs(self): *zip(*input_tokens)) expected_encoded = self.tokenizer.batch_encode_plus( sentence_pairs, - return_tensors='tf', + return_tensors='pt', add_special_tokens=True, padding='longest', truncation='longest_first') @@ -111,7 +111,7 @@ def test_extra_kw(self): **kw) expected_encoded = self.tokenizer.batch_encode_plus( sentences, - return_tensors='tf', + return_tensors='pt', add_special_tokens=True, padding='longest', truncation='longest_first', diff --git a/lit_nlp/examples/glue/models.py b/lit_nlp/examples/glue/models.py index d5f35df8..848de5de 100644 --- a/lit_nlp/examples/glue/models.py +++ b/lit_nlp/examples/glue/models.py @@ -7,8 +7,9 @@ import os import re import threading -from typing import Any, Optional +from typing import Any, Callable, Optional +from absl import logging import attr from lit_nlp.api import model as lit_model from lit_nlp.api import types as lit_types @@ -16,17 +17,13 @@ from lit_nlp.lib import file_cache from lit_nlp.lib import utils import numpy as np -import tensorflow as tf -import tf_keras as keras +import torch import transformers -os.environ["TF_USE_LEGACY_KERAS"] = "1" - JsonDict = lit_types.JsonDict Spec = lit_types.Spec -TFSequenceClassifierOutput = ( - transformers.modeling_tf_outputs.TFSequenceClassifierOutput -) +_MF = transformers.modeling_outputs +SequenceClassifierOutput = _MF.SequenceClassifierOutput @attr.s(auto_attribs=True, kw_only=True) @@ -73,7 +70,7 @@ def init_spec(cls) -> lit_types.Spec: class GlueModel(lit_model.BatchedModel): - """GLUE benchmark model, using Keras/TF2 and Huggingface Transformers. + """GLUE benchmark model, using PyTorch and Huggingface Transformers. This is a general-purpose classification or regression model. It works for one- or two-segment input, and predicts either a multiclass label or @@ -129,14 +126,13 @@ def _load_model(self, model_name_or_path): model_config = transformers.AutoConfig.from_pretrained( model_name_or_path, num_labels=1 if self.is_regression else len(self.config.labels), - return_dict=False, # default for training; overridden for predict output_attentions=self.config.output_attention, ) - self.model = model_utils.load_pretrained( - transformers.TFAutoModelForSequenceClassification, - model_name_or_path, - config=model_config, - ) + cls = transformers.AutoModelForSequenceClassification + self.model = cls.from_pretrained(model_name_or_path, config=model_config) + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.model.to(self.device) + self.model.eval() def _get_tokens(self, ex: JsonDict, field_name: str) -> list[str]: with self._lock: @@ -144,7 +140,9 @@ def _get_tokens(self, ex: JsonDict, field_name: str) -> list[str]: ex[field_name] ) - def _preprocess(self, inputs: Iterable[JsonDict]) -> dict[str, tf.Tensor]: + def _preprocess( + self, inputs: Iterable[JsonDict] + ) -> transformers.BatchEncoding: # Use pretokenized input if available. tokens_a = [self._get_tokens(ex, self.config.text_a_name) for ex in inputs] tokens_b = None @@ -160,27 +158,25 @@ def _preprocess(self, inputs: Iterable[JsonDict]) -> dict[str, tf.Tensor]: tokens_b, max_length=self.config.max_seq_length, ) - return encoded_input # pytype: disable=bad-return-type + encoded_input.to(self.device) + return encoded_input - def _make_dataset(self, inputs: Iterable[JsonDict]) -> tf.data.Dataset: - """Make a tf.data.Dataset from inputs in LIT format.""" - encoded_input = self._preprocess(inputs) + def _make_label_tensor(self, inputs: Iterable[JsonDict]) -> torch.Tensor: + """Make a label tensor from inputs in LIT format.""" if self.is_regression: - labels = tf.constant( - [ex[self.config.label_name] for ex in inputs], dtype=tf.float32 + return torch.tensor( + [ex[self.config.label_name] for ex in inputs], + dtype=torch.float32, + device=self.device, ) else: indexes = [] if self.config.labels is not None: for ex in inputs: indexes.append(self.config.labels.index(ex[self.config.label_name])) - labels = tf.constant( - indexes, - dtype=tf.int64, + return torch.tensor( + indexes, dtype=torch.long, device=self.device ) - # encoded_input is actually a transformers.BatchEncoding - # object, which tf.data.Dataset doesn't like. Convert to a regular dict. - return tf.data.Dataset.from_tensor_slices((dict(encoded_input), labels)) def train( self, @@ -189,44 +185,123 @@ def train( learning_rate=2e-5, batch_size=32, num_epochs=3, - keras_callbacks=None, + on_epoch_end: Optional[Callable[[int, dict[str, float]], None]] = None, ): - """Run fine-tuning.""" - train_dataset = ( - self._make_dataset(train_inputs) - .shuffle(128) - .batch(batch_size) - .repeat(-1) + """Run fine-tuning. + + Args: + train_inputs: list of training examples in LIT format. + validation_inputs: list of validation examples in LIT format. + learning_rate: learning rate for the Adam optimizer. + batch_size: training batch size. + num_epochs: number of epochs to train for. + on_epoch_end: optional callback, called after each epoch with + (epoch_index, logs) where logs contains epoch metrics. + + Returns: + Dict summarizing training history, with keys "epochs", "history", + "params", and "optimizer_params". + """ + self.model.train() + optimizer = torch.optim.Adam( + self.model.parameters(), lr=learning_rate, eps=1e-08 ) + if self.is_regression: + loss_fn = torch.nn.MSELoss() + else: + loss_fn = torch.nn.CrossEntropyLoss() + # Use larger batch for validation since inference is about 1/2 memory usage # of backprop. eval_batch_size = 2 * batch_size - validation_dataset = self._make_dataset(validation_inputs).batch( - eval_batch_size - ) + steps_per_epoch = len(train_inputs) // batch_size + history: dict[str, list[float]] = { + "loss": [], + "val_loss": [], + ("val_rmse" if self.is_regression else "val_accuracy"): [], + } + for epoch in range(num_epochs): + self.model.train() + perm = torch.randperm(len(train_inputs)) + train_inputs = [train_inputs[i] for i in perm.tolist()] + epoch_losses = [] + for step in range(steps_per_epoch): + batch = train_inputs[step * batch_size : (step + 1) * batch_size] + encoded_input = self._preprocess(batch) + labels = self._make_label_tensor(batch) + optimizer.zero_grad() + out: SequenceClassifierOutput = self.model( + **encoded_input, labels=labels + ) + loss = out.loss + loss.backward() + optimizer.step() + epoch_losses.append(loss.item()) - # Prepare model for training. - opt = keras.optimizers.Adam(learning_rate=learning_rate, epsilon=1e-08) + val_loss, val_metric = self._evaluate( + validation_inputs, eval_batch_size, loss_fn + ) + history["loss"].append(float(np.mean(epoch_losses))) + history["val_loss"].append(val_loss) + history[ + ("val_rmse" if self.is_regression else "val_accuracy") + ].append(val_metric) + logs = {k: v[-1] for k, v in history.items()} + logging.info( + "Epoch %d: %s", + epoch, + ", ".join(f"{k}={v:.4f}" for k, v in logs.items()), + ) + if on_epoch_end is not None: + on_epoch_end(epoch, logs) + + self.model.eval() + return { + "epochs": list(range(num_epochs)), + "history": history, + "params": { + "epochs": num_epochs, + "batch_size": batch_size, + "steps_per_epoch": steps_per_epoch, + }, + "optimizer_params": { + "learning_rate": learning_rate, + "epsilon": 1e-08, + }, + } + + def _evaluate( + self, + eval_inputs: list[JsonDict], + batch_size: int, + loss_fn: torch.nn.Module, + ) -> tuple[float, float]: + """Run a validation pass, returning (loss, rmse_or_accuracy).""" + self.model.eval() + losses, sq_errors, correct, total = [], 0.0, 0, 0 + with torch.no_grad(): + for start in range(0, len(eval_inputs), batch_size): + batch = eval_inputs[start : start + batch_size] + encoded_input = self._preprocess(batch) + labels = self._make_label_tensor(batch) + out: SequenceClassifierOutput = self.model( + **encoded_input, labels=labels + ) + losses.append(out.loss.item()) + if self.is_regression: + sq_errors += float(((out.logits[:, 0] - labels) ** 2).sum()) + total += len(batch) + else: + correct += int( + (out.logits.argmax(dim=-1) == labels).sum() + ) + total += len(batch) + val_loss = float(np.mean(losses)) if losses else float("nan") if self.is_regression: - loss = keras.losses.MeanSquaredError() - metric = keras.metrics.RootMeanSquaredError("rmse") + val_metric = float(np.sqrt(sq_errors / max(total, 1))) else: - loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True) - metric = keras.metrics.SparseCategoricalAccuracy("accuracy") - self.model.compile(optimizer=opt, loss=loss, metrics=[metric]) - - steps_per_epoch = len(train_inputs) // batch_size - validation_steps = len(validation_inputs) // eval_batch_size - history = self.model.fit( - train_dataset, - epochs=num_epochs, - steps_per_epoch=steps_per_epoch, - validation_data=validation_dataset, - validation_steps=validation_steps, - callbacks=keras_callbacks, - verbose=2, - ) - return history + val_metric = correct / max(total, 1) + return val_loss, val_metric def save(self, path: str): """Save model weights and tokenizer info. @@ -331,9 +406,9 @@ def _scatter_embs( """Scatters custom passed embeddings into the default model embeddings. Args: - passed_input_embs: [num_scatter_tokens], the custom passed + passed_input_embs: [num_scatter_tokens], the custom passed embeddings to be scattered into the default model embeddings. - input_embs: the default model embeddings. + input_embs: the default model embeddings, as a torch.Tensor or ndarray. batch_indices: the indices of the embeddings to replace in the format (batch_index, sequence_index). offsets: the offset from which to scatter the custom embedding (number of @@ -362,8 +437,17 @@ def _scatter_embs( scatter_indices.append([batch_index, token_index + offset]) # Scatters passed word embeddings into embeddings gathered from tokens. - # [batch_size, num_tokens + num_special_tokens, emb_size] - return tf.tensor_scatter_nd_update(input_embs, scatter_indices, updates) + # [batch_size, num_tokens + num_special_tokens, emb_size] + if not isinstance(input_embs, torch.Tensor): + input_embs = torch.from_numpy(np.asarray(input_embs)) + indices = torch.tensor(scatter_indices, dtype=torch.long).t() + updates_tensor = ( + torch.from_numpy(np.asarray(updates)) + .to(device=input_embs.device, dtype=input_embs.dtype) + ) + return input_embs.index_put( + (indices[0], indices[1]), updates_tensor + ) def scatter_all_embeddings(self, inputs, input_embs): """Scatters custom passed embeddings for text segment inputs. @@ -429,12 +513,12 @@ def get_target_scores(self, inputs: Iterable[JsonDict], scores): Args: inputs: list of input examples - scores: [batch_size, num_classes], either logits or probas + scores: [batch_size, num_classes], either logits or probas Returns: - [batch_size] target scores for each input + [batch_size] target scores for each input """ - arg_max = tf.math.argmax(scores, axis=-1).numpy() + arg_max = scores.argmax(dim=-1).cpu().numpy() grad_classes = [ ex.get(self.config.label_name, arg_max[i]) for (i, ex) in enumerate(inputs) @@ -446,10 +530,12 @@ def get_target_scores(self, inputs: Iterable[JsonDict], scores): grad_idxs.append(self.config.labels.index(label)) else: grad_idxs.append(label) - # list of tuples (batch idx, label idx) - gather_indices = list(enumerate(grad_idxs)) - # [batch_size] - return tf.gather_nd(scores, gather_indices), grad_idxs + # [batch_size] + batch_idx = torch.arange(len(inputs), device=scores.device) + idx_tensor = torch.tensor( + grad_idxs, dtype=torch.long, device=scores.device + ) + return scores[batch_idx, idx_tensor], grad_idxs ## # LIT API implementation @@ -457,48 +543,44 @@ def max_minibatch_size(self): return self.config.inference_batch_size def get_embedding_table(self): - # TODO(b/236276775): Unify on the TFBertEmbeddings.weight API after - # transformers is updated to v4.25.1 (or newer). - if hasattr(self.model.bert.embeddings, "word_embeddings"): - return self.vocab, self.model.bert.embeddings.word_embeddings.numpy() - else: - return self.vocab, self.model.bert.embeddings.weight.numpy() + return ( + self.vocab, + self.model.get_input_embeddings().weight.detach().cpu().numpy(), + ) def predict_minibatch(self, inputs: Iterable[JsonDict]): - # Use watch_accessed_variables to save memory by having the tape do nothing - # if we don't need gradients. - with tf.GradientTape( - watch_accessed_variables=self.config.compute_grads - ) as tape: - encoded_input = self._preprocess(inputs) - - # Gathers word embeddings from BERT model embedding layer using input ids + encoded_input = self._preprocess(inputs) + + grad_context = ( + torch.enable_grad() + if self.config.compute_grads + else torch.no_grad() + ) + with grad_context: + # Gathers word embeddings from the model embedding layer using input ids # of the tokens. input_ids = encoded_input["input_ids"] - word_embeddings = self.model.bert.embeddings.weight - # [batch_size, num_tokens, emb_size] - input_embs = tf.gather(word_embeddings, input_ids) + word_embeddings = self.model.get_input_embeddings().weight + # [batch_size, num_tokens, emb_size] + input_embs = torch.nn.functional.embedding(input_ids, word_embeddings) # Scatter in any passed in embeddings. - # [batch_size, num_tokens, emb_size] + # [batch_size, num_tokens, emb_size] input_embs = self.scatter_all_embeddings(inputs, input_embs) - tape.watch(input_embs) # Watch input_embs for gradient calculation. - - model_inputs = encoded_input.copy() - model_inputs["input_ids"] = None - out: TFSequenceClassifierOutput = self.model( - model_inputs, + model_inputs = { + k: v for k, v in encoded_input.items() if k != "input_ids" + } + out: SequenceClassifierOutput = self.model( + **model_inputs, inputs_embeds=input_embs, - training=False, output_hidden_states=True, output_attentions=True, - return_dict=True, ) batched_outputs = { "input_ids": encoded_input["input_ids"], - "ntok": tf.reduce_sum(encoded_input["attention_mask"], axis=1), + "ntok": encoded_input["attention_mask"].sum(dim=1), "cls_emb": out.hidden_states[-1][:, 0], # last layer, first token } @@ -508,17 +590,17 @@ def predict_minibatch(self, inputs: Iterable[JsonDict]): self._verify_num_layers(out.hidden_states) # [batch_size, num_tokens, 1] - token_mask = tf.expand_dims( - tf.cast(encoded_input["attention_mask"], tf.float32), axis=2 + token_mask = torch.unsqueeze( + encoded_input["attention_mask"].to(torch.float32), dim=2 ) # [batch_size, 1] - denom = tf.reduce_sum(token_mask, axis=1) + denom = token_mask.sum(dim=1) for i, layer_output in enumerate(out.hidden_states): # layer_output is [batch_size, num_tokens, emb_dim] # average over tokens to get [batch_size, emb_dim] batched_outputs[f"layer_{i}/avg_emb"] = ( - tf.reduce_sum(layer_output * token_mask, axis=1) / denom - ) + layer_output * token_mask + ).sum(dim=1) / denom if self.config.output_attention: if len(out.attentions) != self.model.config.num_hidden_layers: @@ -529,37 +611,38 @@ def predict_minibatch(self, inputs: Iterable[JsonDict]): f"{len(out.attentions)}." ) for i, layer_attention in enumerate(out.attentions): + # [batch_size, num_heads, num_tokens, num_tokens] batched_outputs[f"layer_{i+1}/attention"] = layer_attention if self.is_regression: - # [batch_size] - batched_outputs["score"] = tf.squeeze(out.logits, axis=-1) - # [batch_size], a single target per example + # [batch_size] + batched_outputs["score"] = torch.squeeze(out.logits, dim=-1) + # [batch_size], a single target per example scalar_targets = batched_outputs["score"] else: - # [batch_size, num_labels] - batched_outputs["probas"] = tf.nn.softmax(out.logits, axis=-1) - # [batch_size], a single target per example + # [batch_size, num_labels] + batched_outputs["probas"] = torch.softmax(out.logits, dim=-1) + # [batch_size], a single target per example scalar_targets, grad_idxs = self.get_target_scores( inputs, batched_outputs["probas"] ) # TODO(b/294613507): remove once TCAV updated. if self.config.compute_grads: - batched_outputs[self.config.label_name] = tf.convert_to_tensor( - grad_idxs + batched_outputs[self.config.label_name] = torch.tensor( + grad_idxs, device=self.device ) - # Request gradients after the tape is run. + # Request gradients after the forward pass is run. # Note: embs[0] includes position and segment encodings, as well as subword # embeddings. if self.config.compute_grads: - # [batch_size, num_tokens, emb_dim] - batched_outputs["input_emb_grad"] = tape.gradient( - scalar_targets, input_embs - ) + # [batch_size, num_tokens, emb_dim] + batched_outputs["input_emb_grad"] = torch.autograd.grad( + scalar_targets.sum(), input_embs + )[0] detached_outputs = { - k: v.numpy() + k: v.detach().cpu().numpy() for k, v in batched_outputs.items() if v is not None } diff --git a/lit_nlp/examples/notebooks/LIT_components_example.ipynb b/lit_nlp/examples/notebooks/LIT_components_example.ipynb index 6817ac54..f1943774 100644 --- a/lit_nlp/examples/notebooks/LIT_components_example.ipynb +++ b/lit_nlp/examples/notebooks/LIT_components_example.ipynb @@ -139,12 +139,8 @@ }, "outputs": [], "source": [ - "# Fetch the trained model weights and load the model to analyze\n", - "!wget https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz\n", - "!mkdir sst2_tiny\n", - "!tar -xvf sst2_tiny.tar.gz -C sst2_tiny\n", - "\n", - "sentiment_model = models.SST2Model('./sst2_tiny')\n", + "# Load the model to analyze; weights are fetched from HuggingFace on first use\n", + "sentiment_model = models.SST2Model('philschmid/tiny-bert-sst2-distilled')\n", "sentiment_model.input_spec(), sentiment_model.output_spec()" ] }, diff --git a/lit_nlp/examples/notebooks/LIT_sentiment_classifier.ipynb b/lit_nlp/examples/notebooks/LIT_sentiment_classifier.ipynb index a3d9b811..8a2e95b4 100644 --- a/lit_nlp/examples/notebooks/LIT_sentiment_classifier.ipynb +++ b/lit_nlp/examples/notebooks/LIT_sentiment_classifier.ipynb @@ -55,9 +55,7 @@ }, "outputs": [], "source": [ - "# Fetch the trained model weights\n", - "!wget https://storage.googleapis.com/what-if-tool-resources/lit-models/sst2_tiny.tar.gz\n", - "!tar -xvf sst2_tiny.tar.gz" + "# Model weights are fetched from HuggingFace on first use." ] }, { @@ -70,7 +68,7 @@ "source": [ "# Create the LIT widget with the model and dataset to analyze.\n", "datasets = {'sst_dev': data.SST2Data('validation')}\n", - "models = {'sst_tiny': models.SST2Model('./')}\n", + "models = {'sst_tiny': models.SST2Model('philschmid/tiny-bert-sst2-distilled')}\n", "\n", "widget = notebook.LitWidget(models, datasets, port=8890)" ] diff --git a/lit_nlp/examples/penguin/data.py b/lit_nlp/examples/penguin/data.py index 91e3aea3..c32e8470 100644 --- a/lit_nlp/examples/penguin/data.py +++ b/lit_nlp/examples/penguin/data.py @@ -1,13 +1,20 @@ -"""🐧 Penguin tabular dataset from TFDS. +"""🐧 Penguin tabular dataset from the palmerpenguins project. -See https://www.tensorflow.org/datasets/catalog/penguins. for details. +See https://allisonhorst.github.io/palmerpenguins/ for details. """ -from collections.abc import Mapping -from typing import Optional, Union +from typing import Optional + +import pandas as pd + from lit_nlp.api import dataset as lit_dataset from lit_nlp.api import types as lit_types -import tensorflow_datasets as tfds +from lit_nlp.lib import file_cache + +PENGUINS_URL = ( + 'https://raw.githubusercontent.com/allisonhorst/palmerpenguins/main/' + 'inst/extdata/penguins.csv' +) VOCABS = { 'island': ['Biscoe', 'Dream', 'Torgersen'], @@ -28,30 +35,25 @@ class PenguinDataset(lit_dataset.Dataset): """Dataset of penguin tabular data. - From https://www.tensorflow.org/datasets/catalog/penguins. + From https://allisonhorst.github.io/palmerpenguins/. """ - @classmethod - def lit_example_from_record(cls, rec: Mapping[str, Union[float, int]]): - return { - 'body_mass_g': rec['body_mass_g'], - 'culmen_depth_mm': rec['culmen_depth_mm'], - 'culmen_length_mm': rec['culmen_length_mm'], - 'flipper_length_mm': rec['flipper_length_mm'], - 'island': VOCABS['island'][rec['island']], - 'sex': VOCABS['sex'][rec['sex']], - 'species': VOCABS['species'][rec['species']], - } - def __init__(self, max_examples: Optional[int] = None): - peng = tfds.load('penguins/simple', download=True, try_gcs=True) - dataset_df = tfds.as_dataframe(peng['train']) + path = file_cache.cached_path(PENGUINS_URL) + dataset_df = pd.read_csv(path) + # Match the field names used by the LIT spec. + dataset_df = dataset_df.rename(columns={ + 'bill_length_mm': 'culmen_length_mm', + 'bill_depth_mm': 'culmen_depth_mm', + }) + dataset_df['sex'] = dataset_df['sex'].str.capitalize() - # Filter out invalid rows. - dataset_df = dataset_df.loc[dataset_df['sex'] != 2] + # Filter out rows with missing values. + fields = list(self.spec().keys()) + dataset_df = dataset_df.dropna(subset=fields) records = dataset_df.to_dict(orient='records') self._examples = [ - PenguinDataset.lit_example_from_record(rec) for rec in records + {field: rec[field] for field in fields} for rec in records ][:max_examples] @classmethod diff --git a/lit_nlp/examples/penguin/demo.py b/lit_nlp/examples/penguin/demo.py index b6d1e855..c0344580 100644 --- a/lit_nlp/examples/penguin/demo.py +++ b/lit_nlp/examples/penguin/demo.py @@ -20,12 +20,8 @@ from lit_nlp.examples.penguin import data as penguin_data from lit_nlp.examples.penguin import model as penguin_model -MODEL_PATH = 'https://storage.googleapis.com/what-if-tool-resources/lit-models/penguin.h5' # pylint: disable=line-too-long - FLAGS = flags.FLAGS FLAGS.set_default('default_layout', 'penguins') -_MODEL_PATH = flags.DEFINE_string('model_path', MODEL_PATH, - 'Path to load trained model.') _MAX_EXAMPLES = flags.DEFINE_integer( 'max_examples', @@ -69,7 +65,7 @@ def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]: if len(argv) > 1: raise app.UsageError('Too many command-line arguments.') - models = {'species classifier': penguin_model.PenguinModel(_MODEL_PATH.value)} + models = {'species classifier': penguin_model.PenguinModel()} datasets = {'penguins': penguin_data.PenguinDataset()} # Truncate datasets if --max_examples is set. if _MAX_EXAMPLES.value is not None: diff --git a/lit_nlp/examples/penguin/model.py b/lit_nlp/examples/penguin/model.py index 510e4910..7e852c29 100644 --- a/lit_nlp/examples/penguin/model.py +++ b/lit_nlp/examples/penguin/model.py @@ -1,72 +1,62 @@ -"""🐧 TensorFlow Keras model for the Penguin dataset.""" - -import os +"""🐧 Scikit-learn model for the Penguin dataset.""" from lit_nlp.api import model as lit_model from lit_nlp.api import types as lit_types from lit_nlp.examples.penguin import data as penguin_data -from lit_nlp.lib import file_cache import numpy as np -import tf_keras as keras - -os.environ['TF_USE_LEGACY_KERAS'] = '1' +from sklearn.linear_model import LogisticRegression +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler _VOCABS = penguin_data.VOCABS class PenguinModel(lit_model.BatchedModel): - """TensorFlow Keras model for penguin classification.""" - - def __init__(self, path: str): - path = file_cache.cached_path(path) - self.model = keras.models.load_model(path) - # Feature column means and variance to normalize values before - # prediction. - self.means = np.array([ - 4.23795547e+03, 1.71222672e+01, 4.40004048e+01, 2.01587045e+02, - 5.14170040e-01, 3.40080972e-01, 1.45748988e-01, 4.89878543e-01, - 5.10121457e-01 - ]) - self.vars = np.array([ - 6.29523451e+05, 4.17662973e+00, 2.76167220e+01, 1.97108820e+02, - 2.49799210e-01, 2.24425904e-01, 1.24506220e-01, 2.49897556e-01, - 2.49897556e-01 + """Scikit-learn model for penguin classification. + + The model is trained on the penguin dataset at construction time; there + is no external model file to download. + """ + + def __init__(self): + examples = penguin_data.PenguinDataset().examples + features = np.array([self._convert_input(ex) for ex in examples]) + labels = [ex['species'] for ex in examples] + self.model = make_pipeline( + StandardScaler(), + LogisticRegression(max_iter=1000), + ) + self.model.fit(features, labels) + + def _convert_input(self, inp): + ex = np.array([ + inp['body_mass_g'], inp['culmen_depth_mm'], inp['culmen_length_mm'], + inp['flipper_length_mm'], 0, 0, 0, 0, 0 ]) + # Set one-hot encodings of categorical features. + island_index = _VOCABS['island'].index(inp['island']) + sex_index = _VOCABS['sex'].index(inp['sex']) + # Island one-hot encodings start at input index 4. + ex[island_index + 4] = 1 + # Sex one-hot encodings start at input index 7. + ex[sex_index + 7] = 1 + return ex def max_minibatch_size(self) -> int: return 32 def predict_minibatch(self, inputs): + adjusted_inputs = np.array([self._convert_input(inp) for inp in inputs]) + model_output = self.model.predict_proba(adjusted_inputs) - def convert_input(inp): - ex = np.array([ - inp['body_mass_g'], inp['culmen_depth_mm'], inp['culmen_length_mm'], - inp['flipper_length_mm'], 0, 0, 0, 0, 0 - ]) - # Set one-hot encodings of categorical features. - island_index = _VOCABS['island'].index(inp['island']) - sex_index = _VOCABS['sex'].index(inp['sex']) - # Island one-hot encodings start at input index 4. - ex[island_index + 4] = 1 - # Sex one-hot encodings start at input index 7. - ex[sex_index + 7] = 1 - - # Normalize feature values. - ex = (ex - self.means) / (self.vars**0.5) - return ex.tolist() - - adjusted_inputs = [convert_input(inp) for inp in inputs] - model_output = self.model.predict(adjusted_inputs) - + # sklearn sorts classes alphabetically, which matches VOCABS['species']. + assert list(self.model.classes_) == _VOCABS['species'] ret = [{'predicted_species': out} for out in model_output] return ret @classmethod def init_spec(cls) -> lit_types.Spec: - - default_model_path = ('https://storage.googleapis.com/' - 'what-if-tool-resources/lit-models/penguin.h5') - return {'path': lit_types.String(default=default_model_path)} + return {} def input_spec(self) -> lit_types.Spec: return penguin_data.INPUT_SPEC diff --git a/lit_nlp/examples/penguin/model_int_test.py b/lit_nlp/examples/penguin/model_int_test.py index c9dca404..ab6dc8c0 100644 --- a/lit_nlp/examples/penguin/model_int_test.py +++ b/lit_nlp/examples/penguin/model_int_test.py @@ -8,9 +8,8 @@ class PenguinModelIntTest(absltest.TestCase): """Test that model class can predict.""" def test_model(self): - # Create model. - model_path = "https://storage.googleapis.com/what-if-tool-resources/lit-models/penguin.h5" # pylint: disable=line-too-long - model = penguin_model.PenguinModel(model_path) + # Create model. The model trains on the penguin dataset at init. + model = penguin_model.PenguinModel() # Run prediction to ensure no failure. model_in = [{ diff --git a/lit_nlp/examples/prompt_debugging/keras_lms.py b/lit_nlp/examples/prompt_debugging/keras_lms.py deleted file mode 100644 index 84755583..00000000 --- a/lit_nlp/examples/prompt_debugging/keras_lms.py +++ /dev/null @@ -1,536 +0,0 @@ -"""LIT model wrappers for generic instrumented Keras LMs.""" - -from collections.abc import Sequence -import functools -import inspect -from typing import Optional - -from absl import logging -import keras -from keras_nlp import models as keras_models -from lit_nlp.api import model as lit_model -from lit_nlp.api import types as lit_types -from lit_nlp.examples.prompt_debugging import constants as pd_constants -from lit_nlp.examples.prompt_debugging import utils as pd_utils -from lit_nlp.lib import file_cache -from lit_nlp.lib import utils as lit_utils - - -# pylint: disable=g-import-not-at-top -# pytype: disable=import-error -# NOTE: The Keras backend must be set before loading the Keras library. You can -# set the backend using the KERAS_BACKEND environment variable or your -# ~/.keras/keras.json configuration file. For more information, see: -# https://keras.io/getting_started/#configuring-your-backend -if keras.backend.backend() == "tensorflow": - import tensorflow as tf -elif keras.backend.backend() == "torch": - import torch -else: - # TODO(b/333373960): Update imports once a JAX salience is supported. - raise ValueError(f"Unsupported backend: {keras.backend.backend()}") -# pytype: enable=import-error -# pylint: enable=g-import-not-at-top - - -_DEFAULT_MAX_LENGTH = 1024 - - -class _KerasBaseModel(lit_model.BatchedModel): - """Base LIT model wrapper class for Keras on TensorFlow.""" - - # TODO(lit-dev): pytype annotations for model= ? - # Should be keras_nlp.models.generative_task.GenerativeTask - def __init__( - self, - model: Optional[keras_models.CausalLM] = None, - model_name_or_path: Optional[str] = None, - max_length: int = _DEFAULT_MAX_LENGTH, - dynamic_sequence_length: bool = True, - batch_size: int = 16, - ): - """Base wrapper for a Keras CausalLM supporting the layer_intercept_fn API. - - Model should support the following methods: - - .generate() - - .score()* - - .preprocessor.generate_preprocess() - . .preprocessor.tokenizer.id_to_token() - . .backbone.token_embedding() - - * The score function should accept layer_intercept_fn= as a way to intercept - and manipulate activations between layers. We use this for salience, below. - - Args: - model: A pre-loaded Keras CausalLM, prioritized over model_name_or_path. - model_name_or_path: A URL, path, or preset name for the model to load, - max_length: max sequence length - dynamic_sequence_length: if true, will trim padding to the length of the - longest sequence in a batch. Recommended for CPU and GPU usage, but may - be disabled for compilation where a fixed shape is required. - batch_size: batch size - """ - super().__init__() - - if model is not None: - self.model = model - elif model_name_or_path is not None: - if ( - is_tar_gz := model_name_or_path.endswith(".tar.gz") - ) or file_cache.is_remote(model_name_or_path): - model_name_or_path = file_cache.cached_path( - model_name_or_path, - extract_compressed_file=is_tar_gz, - ) - self.model = keras_models.CausalLM.from_preset(model_name_or_path) - else: - raise ValueError("Must provide either model or model_name_or_path.") - - self.batch_size = batch_size - self.max_length = max_length - self.dynamic_sequence_length = dynamic_sequence_length - - # map ids: [batch_size, num_tokens] - # to embs: [batch_size, num_tokens, emb_dim] - self.embedder = self.model.backbone.token_embedding - - def encode_inputs(self, texts: Sequence[str]): - """Encode inputs, with optional dynamic trimming. - - By default, the model's generate_preprocess() pads to a fixed sequence - length, either specified as sequence_length= or using an internal default. - - Here, we optionally trim this to remove extraneous padding positions based - on the actual contents of the minibatch. This can greatly speed up - performance when running on CPU or GPU. - - Args: - texts: list of input strings - - Returns: - A dict[str, Tensor] compatible with model.score(), etc. functions. - """ - # First: pack to max_length - encoded_inputs = self.model.preprocessor.generate_preprocess( - texts, sequence_length=self.max_length - ) - if not self.dynamic_sequence_length: - return encoded_inputs - - # Trim to the maximum length needed to contain any non-padding tokens. - mask = encoded_inputs["padding_mask"] - - if keras.backend.backend() == "tensorflow": - max_indices = [tf.reduce_max(tf.where(row)) for row in mask] - elif keras.backend.backend() == "torch": - max_indices = [torch.max(torch.where(row)[0]) for row in mask] - else: - raise ValueError(f"Unsupported backend: {keras.backend.backend()}") - # Find position of last 'True' in each row. - seq_ends: Sequence[int] = [ - keras.ops.convert_to_numpy(i).tolist() + 1 for i in max_indices - ] - longest_sequence = max(seq_ends) - # TODO(lit-dev): remove this line, or make it logging.debug ? - logging.info( - "Trimming batch to trimmed_length = %d based on sequence ends %s", - longest_sequence, - seq_ends, - ) - # Actually trim the input tensors. - return {k: v[:, :longest_sequence] for k, v in encoded_inputs.items()} - - def clean_subword_token(self, tok: str) -> str: - """Clean up special subword token from the tokenizers if necessary. - - Args: - tok: the token to clean up. - Returns: - The replaced token if the provided token matches the special subword token - below; otherwise, the original token is returned. - """ - # For GPT2 tokenizer. - tok = tok.replace("Ċ", "\n") # newlines - tok = tok.replace("Ġ", "▁") # start of word -> magic underscore - # For SentencePiece Tokenizer. - tok = tok.replace("<0x0A>", "\n") # newlines - return tok - - def ids_to_clean_tokens(self, ids: Sequence[int]) -> Sequence[str]: - return [ - self.clean_subword_token( - self.model.preprocessor.tokenizer.id_to_token(id) - ) - for id in ids - ] - - @classmethod - def from_loaded(cls, existing: "_KerasBaseModel", *args, **kw): - """Share weights and underlying Keras model with another instance.""" - return cls(model=existing.model, *args, **kw) - - def max_minibatch_size(self) -> int: - return self.batch_size - - @classmethod - def init_spec(cls): - # Cannot initialize from spec, because we need a Keras model object. - return None - - def input_spec(self): - return pd_constants.INPUT_SPEC - - -class KerasGenerationModel(_KerasBaseModel): - """LIT model wrapper for generating text with Keras on TensorFlow. - - This class accepts a loaded model and provides the LIT-required functions plus - additional helper functions for generation tasks. - - This class supports generation and pass-through modes. If a dataset provides a - pre-populated 'response' column then this model will return that text instead - of generating new text from the 'prompt'. This allows the same model wrapper - to be efficiently used to examine saved results from bulk-inference pipelines - and new generations from, e.g., counterfactually generated examples, or novel - evaluation datasets. - """ - - def __init__(self, *args, output_embeddings=True, **kw): - super().__init__(*args, **kw) - self.output_embeddings = output_embeddings - - def embed_texts(self, texts: Sequence[str]): - processed_inputs = self.encode_inputs(texts) - # [batch_size, num_tokens, emb_dim] - embs = self.embedder(processed_inputs["token_ids"]) - # [batch_size, num_tokens] - mask = processed_inputs["padding_mask"] - return embs, mask - - def embed_and_mean_pool(self, texts: Sequence[str]): - """Return a single vector for each text.""" - embs, mask = self.embed_texts(texts) - # [batch_size, num_tokens, 1] - cast_mask = keras.ops.cast(mask, dtype=embs.dtype) - - if keras.backend.backend() == "tensorflow": - expanded_mask = tf.expand_dims(cast_mask, axis=2) - pooled_embs = tf.reduce_sum( - expanded_mask * embs, axis=1, keepdims=True - ) / tf.reduce_sum(expanded_mask, axis=1, keepdims=True) - return tf.squeeze(pooled_embs, axis=1) - elif keras.backend.backend() == "torch": - expanded_mask = torch.unsqueeze(cast_mask, dim=2) - pooled_embs = torch.sum( - expanded_mask * embs, dim=1, keepdim=True - ) / torch.sum(expanded_mask, dim=1, keepdim=True) - return torch.squeeze(pooled_embs, dim=1) - else: - raise ValueError(f"Unsupported backend: {keras.backend.backend()}") - - def predict_minibatch( - self, - inputs: list[lit_types.JsonDict], - ) -> list[lit_types.JsonDict]: - prompts: Sequence[str] = [ - ex[pd_constants.FieldNames.PROMPT] for ex in inputs - ] - - # TODO(lit-dev): suppport loading cached responses here, since running - # generation can be expensive. - full_responses: Sequence[str] = list( - self.model.generate(prompts, max_length=self.max_length) - ) - # Model outputs include the prompt, so trim that off and just return the - # generated portion. - responses: Sequence[str] = [ - response[len(prompt) :] - for response, prompt in zip(full_responses, prompts) - ] - - outputs = [ - {pd_constants.FieldNames.RESPONSE: response} for response in responses - ] - - if self.output_embeddings: - prompt_embeddings = self.embed_and_mean_pool(prompts) - # TODO(lit-dev): embed prompt + response and trim embedding instead? - # Or just embed full_response. - response_embeddings = self.embed_and_mean_pool(responses) - - for o, p, r in zip(outputs, prompt_embeddings, response_embeddings): - o[pd_constants.FieldNames.PROMPT_EMBEDDINGS] = ( - keras.ops.convert_to_numpy(p) - ) - o[pd_constants.FieldNames.RESPONSE_EMBEDDINGS] = ( - keras.ops.convert_to_numpy(r) - ) - - return outputs - - def output_spec(self) -> lit_types.Spec: - ret = pd_constants.OUTPUT_SPEC_GENERATION - if self.output_embeddings: - return ret | pd_constants.OUTPUT_SPEC_GENERATION_EMBEDDINGS - return ret - - -class KerasSalienceModel(_KerasBaseModel): - """LIT model wrapper for computing salience with Keras on TensorFlow. - - This class accepts a loaded model and provides the LIT-required functions plus - additional helper functions to convert and clean tokens and to compute - sequence salience. - - This class does not support generation; use the KerasGenerationModel class to - generate the text for which this class will compute salience. - """ - - def __init__(self, *args, **kw): - super().__init__(*args, **kw) - - score_fn = getattr(self.model, "score", None) - - if score_fn is None or not inspect.ismethod(score_fn): - raise TypeError( - "Salience is computed via a .score() API, which is not supported by " - "all KerasNLP CausalLM models. Please provide a model that " - "supports this API." - ) - - def _pred(self, input_ids, padding_mask, target_masks): - """Predict a batch of tokenized text. - - Args: - input_ids: A Tensor with shape [batch_size, num_tokens] - padding_mask: A Tensor with shape [batch_size, num_tokens] - target_masks: A Numpy Array with shape [batch_size, num_tokens] - - Returns: - Batched outputs for post-processing. - """ - ## - # Process target masks - # It doesn't make sense to interpret the first token, since it is not ever - # predicted. But we need to ensure that the mask[0] is zero, so it doesn't - # cause problems when 'rolled' to the last position below. - seq_len = keras.ops.shape(input_ids)[1] - pad_fn = functools.partial( - lit_utils.pad1d, - min_len=seq_len, - max_len=seq_len, - pad_val=0, - pad_left=False, - ) - - modified_masks = [[0] + list(mask[1:]) for mask in target_masks] - stacked_padded_masks = keras.ops.stack( - [pad_fn(mask) for mask in modified_masks], - axis=0, - ) - # Shift masks back so they align with the target_ids generated in the - # backend-specific prediction functions. - rolled_masks = keras.ops.roll(stacked_padded_masks, shift=-1, axis=1) - loss_mask = keras.ops.convert_to_tensor(rolled_masks, dtype="bool") - - pred_kw_args = { - "input_ids": input_ids, - "padding_mask": padding_mask, - "loss_mask": loss_mask, - } - if keras.backend.backend() == "tensorflow": - grad_l2, grad_dot_input = self._pred_tf(**pred_kw_args) - elif keras.backend.backend() == "jax": - grad_l2, grad_dot_input = self._pred_jax(**pred_kw_args) - elif keras.backend.backend() == "torch": - grad_l2, grad_dot_input = self._pred_torch(**pred_kw_args) - else: - raise ValueError(f"Unsupported backend: {keras.backend.backend()}") - - batched_outputs = { - "input_ids": input_ids, - "padding_mask": padding_mask, - pd_constants.FieldNames.GRAD_NORM: grad_l2, - pd_constants.FieldNames.GRAD_DOT_INPUT: grad_dot_input, - } - - return batched_outputs - - def _pred_tf(self, input_ids, padding_mask, loss_mask): - # [batch_size, num_tokens]; ignore the last one in each row. - target_ids = tf.roll(input_ids, shift=-1, axis=1) - embeddings = None - - with tf.GradientTape(watch_accessed_variables=False) as tape: - - def layer_intercept_fn(x, i): - if i == -1: - nonlocal embeddings, tape - embeddings = x - tape.watch(embeddings) - return x - - # [batch_size, num_tokens] - per_token_loss = self.model.score( - token_ids=input_ids, - padding_mask=padding_mask, - scoring_mode="loss", - layer_intercept_fn=layer_intercept_fn, - target_ids=target_ids, - ) - masked_loss = per_token_loss * keras.ops.cast( - loss_mask, per_token_loss.dtype - ) - - # [batch_size, num_tokens, hdim] - grads = tape.gradient(masked_loss, embeddings) - # [batch_size, num_tokens] - grad_l2 = tf.norm(grads, axis=2) - # [batch_size, num_tokens] - grad_dot_input = tf.reduce_sum(grads * embeddings, axis=2) - return grad_l2, grad_dot_input - - # TODO(b/333373960): Implement salience computation for JAX. - def _pred_jax(self, input_ids, padding_mask, loss_mask): - # NOTE: JAX computes gradients automatically w.r.t function inputs and - # outputs. The score function takes token_ids as its input but salience is - # computed w.r.t. the embeddings, thus JAX cannot differentiate the loss - # w.r.t. the embeddings and taking gradients w.r.t. the token_ids is not - # equivalent. For now, we raise an error if using JAX. - raise NotImplementedError("JAX backend not supported for salience.") - - def _pred_torch(self, input_ids, padding_mask, loss_mask): - target_ids = torch.roll(input_ids, shifts=-1, dims=1) - embeddings = None - - def layer_intercept_fn(x, i): - if i == -1: - nonlocal embeddings - embeddings = x - return x - - per_token_loss = self.model.score( - token_ids=input_ids, - padding_mask=padding_mask, - scoring_mode="loss", - layer_intercept_fn=layer_intercept_fn, - target_ids=target_ids, - ) - - if embeddings is None: - raise ValueError("Embeddings are None after scoring.") - - masked_loss = per_token_loss * keras.ops.cast( - loss_mask, per_token_loss.dtype - ) - - # [batch_size, num_tokens, hdim] - grads = torch.autograd.grad( - masked_loss, embeddings, grad_outputs=torch.ones_like(masked_loss) - )[0] - embeddings = embeddings.detach() - # [batch_size, num_tokens] - grad_l2 = torch.norm(grads, dim=2) - # [batch_size, num_tokens] - grad_dot_input = torch.sum(grads * embeddings, dim=2) - return grad_l2, grad_dot_input - - def _postprocess(self, preds): - """Post-process single-example preds. Operates on numpy arrays.""" - mask = preds.pop("padding_mask").astype(bool) - ids = preds.pop("input_ids")[mask] - preds[pd_constants.FieldNames.TOKENS] = self.ids_to_clean_tokens(ids) - for key in lit_utils.find_spec_keys( - self.output_spec(), lit_types.TokenScores - ): - preds[key] = preds[key][mask] - # First token () is not actually predicted, so return 0 for loss. - # preds[FieldNames.TOKEN_LOSS][0] = 0 - - return preds - - def predict_minibatch(self, inputs): - """Predict on a single minibatch of examples.""" - texts: Sequence[str] = [ - ex[pd_constants.FieldNames.PROMPT] - + ex.get(pd_constants.FieldNames.TARGET, "") - for ex in inputs - ] - preprocessed_texts = self.encode_inputs(texts) - sequence_ids = preprocessed_texts["token_ids"] - padding_mask = preprocessed_texts["padding_mask"] - - target_masks = [ - ex.get(pd_constants.FieldNames.TARGET_MASK, []) for ex in inputs - ] - - # Get the predictions. - batched_outputs = self._pred(sequence_ids, padding_mask, target_masks) - # Convert to numpy for post-processing. - detached_outputs = { - k: keras.ops.convert_to_numpy(v) for k, v in batched_outputs.items() - } - # Split up batched outputs, then post-process each example. - unbatched_outputs = lit_utils.unbatch_preds(detached_outputs) - return map(self._postprocess, unbatched_outputs) - - def input_spec(self): - return super().input_spec() | pd_constants.INPUT_SPEC_SALIENCE - - def output_spec(self) -> lit_types.Spec: - return pd_constants.OUTPUT_SPEC_SALIENCE - - -class KerasTokenizerModel(_KerasBaseModel): - """LIT model wrapper for tokenizing text with Keras on TensorFlow. - - This class accepts a loaded model and provides the LIT-required functions plus - additional helper functions to convert and clean tokens. - """ - - def _postprocess(self, preds): - """Post-process single-example preds. Operates on numpy arrays.""" - # Be sure to cast to bool, otherwise this will select intger positions 0, 1 - # rather than acting as a boolean mask. - mask = preds.pop("padding_mask").astype(bool) - ids = preds.pop("token_ids")[mask] - preds[pd_constants.FieldNames.TOKENS] = self.ids_to_clean_tokens(ids) - return preds - - def predict_minibatch(self, inputs): - """Tokenize a single minibatch of examples.""" - texts: Sequence[str] = [ - ex[pd_constants.FieldNames.PROMPT] - + ex.get(pd_constants.FieldNames.TARGET, "") - for ex in inputs - ] - preprocessed_texts = self.encode_inputs(texts) - batched_outputs = { - "token_ids": preprocessed_texts["token_ids"], - "padding_mask": preprocessed_texts["padding_mask"], - } - # Convert to numpy for post-processing. - detached_outputs = { - k: keras.ops.convert_to_numpy(v) for k, v in batched_outputs.items() - } - # Split up batched outputs, then post-process each example. - unbatched_outputs = lit_utils.unbatch_preds(detached_outputs) - return map(self._postprocess, unbatched_outputs) - - def output_spec(self) -> lit_types.Spec: - return pd_constants.OUTPUT_SPEC_TOKENIZER - - -def initialize_model_group_for_salience( - new_name: str, **kw -) -> lit_model.ModelMap: - """Creates '{name}' and '_{name}_salience' and '_{name}_tokenizer'.""" - salience_name, tokenizer_name = pd_utils.generate_model_group_names(new_name) - generation_model = KerasGenerationModel(**kw) - salience_model = KerasSalienceModel(model=generation_model.model, **kw) - tokenizer_model = KerasTokenizerModel(model=generation_model.model, **kw) - return { - new_name: generation_model, - salience_name: salience_model, - tokenizer_name: tokenizer_model, - } diff --git a/lit_nlp/examples/prompt_debugging/models.py b/lit_nlp/examples/prompt_debugging/models.py index ca9c797e..e1705daf 100644 --- a/lit_nlp/examples/prompt_debugging/models.py +++ b/lit_nlp/examples/prompt_debugging/models.py @@ -1,7 +1,6 @@ """Methods for configuring models for prompt debugging.""" from collections.abc import Sequence -import os from typing import Optional from absl import logging @@ -11,41 +10,34 @@ DEFAULT_BATCH_SIZE = 1 -DEFAULT_DL_FRAMEWORK = "kerasnlp" DEFAULT_DL_RUNTIME = "torch" -DEFAULT_MODELS = ["gemma_1.1_instruct_2b_en:gemma_1.1_instruct_2b_en"] +DEFAULT_MODELS = ["gemma-2b-it:google/gemma-2b-it"] DEFAULT_PRECISION = "bfloat16" DEFAULT_SEQUENCE_LENGTH = 512 def _initialize_modeling_environment( - dl_framework: str, dl_runtime: str, precision: str, ) -> None: """Configure the modeling environment.""" - if dl_framework == "kerasnlp": - # NOTE: Keras requires that the KERAS_BACKEND variable is set before import. - os.environ["KERAS_BACKEND"] = dl_runtime - - # NOTE: Imported here and not at the top of the file to avoid - # initialization issues with the environment variables above. - import keras # pylint: disable=g-import-not-at-top - - keras.config.set_floatx(precision) - elif dl_runtime == "torch": - # NOTE: Keras sets precision for all backends with set_floatx(), but for - # HuggingFace Transformers with PyTorch we need to set it explicitly. - import torch # pylint: disable=g-import-not-at-top # pytype: disable=import-error - - torch.set_default_dtype( - torch.bfloat16 if precision == "bfloat16" else torch.float32 + if dl_runtime != DEFAULT_DL_RUNTIME: + raise ValueError( + f"The provided value `{dl_runtime}` for dl_runtime is not supported," + f" only `{DEFAULT_DL_RUNTIME}` is available." ) + # NOTE: Imported here and not at the top of the file to avoid + # initialization issues with the environment variables above. + import torch # pylint: disable=g-import-not-at-top # pytype: disable=import-error + + torch.set_default_dtype( + torch.bfloat16 if precision == "bfloat16" else torch.float32 + ) + def get_models( models_config: Optional[Sequence[str]] = None, - dl_framework: str = DEFAULT_DL_FRAMEWORK, dl_runtime: str = DEFAULT_DL_RUNTIME, precision: str = DEFAULT_PRECISION, batch_size: int = DEFAULT_BATCH_SIZE, @@ -55,12 +47,10 @@ def get_models( Args: models_config: A list of model names and paths to load from, as - "model:path", where path can be a URL, a local file path, or the name of a - preset for the configured deep learning framework. - dl_framework: The deep learning framework that loads and runs the model on - the runtime, `models_config.path` incompatibilities will result in errors. - dl_runtime: The deep learning runtime that the model runs on, either - "tensorflow" or "torch". All loaded models will use the same runtime, + "model:path", where path can be a URL, a local file path, or the name of + a HuggingFace Transformers model. + dl_runtime: The deep learning runtime that the model runs on, only "torch" + is supported. All loaded models will use the same runtime, incompatibilities will result in errors. precision: Floating point precision for the models, either `bfloat16` or `float32`. @@ -76,7 +66,9 @@ def get_models( # NOTE: Always call this function before initializing models to ensure the # environment is properly configured. - _initialize_modeling_environment(dl_framework, dl_runtime, precision) + _initialize_modeling_environment(dl_runtime, precision) + + from lit_nlp.examples.prompt_debugging import transformers_lms # pylint: disable=g-import-not-at-top # pytype: disable=import-error models: dict[str, lit_model.Model] = {} for model_string in models_config: @@ -84,44 +76,28 @@ def get_models( model_name, path = model_string.split(":", 1) logging.info("Loading model '%s' from '%s'", model_name, path) - if dl_framework == "kerasnlp": - from lit_nlp.examples.prompt_debugging import keras_lms # pylint: disable=g-import-not-at-top # pytype: disable=import-error - - models |= keras_lms.initialize_model_group_for_salience( - model_name, - model_name_or_path=path, - max_length=max_length, - batch_size=batch_size, - ) - else: - from lit_nlp.examples.prompt_debugging import transformers_lms # pylint: disable=g-import-not-at-top # pytype: disable=import-error - - models |= transformers_lms.initialize_model_group_for_salience( - model_name, - model_name_or_path=path, - batch_size=batch_size, - framework=dl_runtime, - max_length=max_length, - ) + models |= transformers_lms.initialize_model_group_for_salience( + model_name, + model_name_or_path=path, + batch_size=batch_size, + framework=dl_runtime, + max_length=max_length, + ) return models def get_model_loaders( - dl_framework: str = DEFAULT_DL_FRAMEWORK, dl_runtime: str = DEFAULT_DL_RUNTIME, batch_size: int = DEFAULT_BATCH_SIZE, max_length: int = DEFAULT_SEQUENCE_LENGTH, ) -> lit_app.ModelLoadersMap: - """Get the model loader for the configured framework and runtime. + """Get the model loader for the configured runtime. Args: - dl_framework: The deep learning framework that loads and runs the model on - the runtime, all models are loaded with the same framework, + dl_runtime: The deep learning runtime that the model runs on, only "torch" + is supported. All models are loaded with the same runtime, `model_name_or_path` incompatibilities will result in errors. - dl_runtime: The deep learning runtime that the model runs on, either - "tensorflow" or "torch". All loaded models will use the same runtime, - incompatibilities will result in errors. batch_size: The default batch size. max_length: The default maximum sequence length. @@ -129,7 +105,9 @@ def get_model_loaders( A mapping from model name to initialized LIT model. """ - common_init_spec: lit_types.Spec = { + from lit_nlp.examples.prompt_debugging import transformers_lms # pylint: disable=g-import-not-at-top # pytype: disable=import-error + + transformers_init_spec: lit_types.Spec = { "model_name_or_path": lit_types.String(), "batch_size": lit_types.Integer( default=batch_size, min_val=1, max_val=64, required=False @@ -137,37 +115,14 @@ def get_model_loaders( "max_length": lit_types.Integer( default=max_length, min_val=1, max_val=2048, required=False ), + "framework": lit_types.CategoryLabel( + vocab=transformers_lms.SUPPORTED_ML_RUNTIMES, default=dl_runtime + ), } - if dl_framework == "kerasnlp": - from lit_nlp.examples.prompt_debugging import keras_lms # pylint: disable=g-import-not-at-top # pytype: disable=import-error - - keras_init_spec: lit_types.Spec = { - **common_init_spec, - "dynamic_sequence_length": lit_types.Boolean( - default=True, required=False - ), - } - - return { - "Keras LLM": ( - keras_lms.initialize_model_group_for_salience, - keras_init_spec, - ) - } - else: - from lit_nlp.examples.prompt_debugging import transformers_lms # pylint: disable=g-import-not-at-top # pytype: disable=import-error - - transformers_init_spec: lit_types.Spec = { - **common_init_spec, - "framework": lit_types.CategoryLabel( - vocab=transformers_lms.SUPPORTED_ML_RUNTIMES, default=dl_runtime - ), - } - - return { - "Transformers LLM": ( - transformers_lms.initialize_model_group_for_salience, - transformers_init_spec, - ) - } + return { + "Transformers LLM": ( + transformers_lms.initialize_model_group_for_salience, + transformers_init_spec, + ) + } diff --git a/lit_nlp/examples/prompt_debugging/notebook.py b/lit_nlp/examples/prompt_debugging/notebook.py index df4d5bea..84251693 100644 --- a/lit_nlp/examples/prompt_debugging/notebook.py +++ b/lit_nlp/examples/prompt_debugging/notebook.py @@ -14,8 +14,6 @@ def make_notebook_widget( *, # keep-sorted start batch_size: int = models.DEFAULT_BATCH_SIZE, - dl_framework: str = "kerasnlp", - dl_runtime: str = "tensorflow", max_examples: int = datasets.DEFAULT_MAX_EXAMPLES, precision: str = "bfloat16", # keep-sorted end, @@ -31,11 +29,6 @@ def make_notebook_widget( "model:path", where path can be a URL, a local file path, or the name of a preset for the configured deep learning framework. batch_size: The number of examples the model will process per batch. - dl_framework: The deep learning framework that loads and runs the model on - the runtime, `models_config.path` incompatibilities will result in errors. - dl_runtime: The deep learning runtime that the model runs on, either - "tensorflow" or "torch". All loaded models will use the same runtime, - incompatibilities will result in errors. max_examples: Maximum number of examples in each loaded dataset. precision: Floating point precision for the models, either `bfloat16` or `float32`. @@ -49,8 +42,6 @@ def make_notebook_widget( return lit_notebook.LitWidget( models=models.get_models( models_config=models_config, - dl_framework=dl_framework, - dl_runtime=dl_runtime, precision=precision, batch_size=batch_size, ), @@ -60,8 +51,6 @@ def make_notebook_widget( layouts=layouts.PROMPT_DEBUGGING_LAYOUTS, default_layout=layouts.LEFT_RIGHT, model_loaders=models.get_model_loaders( - dl_framework=dl_framework, - dl_runtime=dl_runtime, batch_size=batch_size, max_length=models.DEFAULT_SEQUENCE_LENGTH, ), diff --git a/lit_nlp/examples/prompt_debugging/server.py b/lit_nlp/examples/prompt_debugging/server.py index 90104708..651c2bb1 100644 --- a/lit_nlp/examples/prompt_debugging/server.py +++ b/lit_nlp/examples/prompt_debugging/server.py @@ -1,14 +1,14 @@ r"""Server for sequence salience with a left-to-right language model. To use with the Gemma, Llama, or Mistral models, install the latest versions of -Keras, KerasNLP, and/or HuggingFace Transformers: +HuggingFace Transformers: - pip install keras>=3.1.0 keras-nlp>=0.9.0 transformers>=4.38.0 + pip install transformers>=4.38.0 -To run with the default configuration (Gemma on TensorFlow via Keras): +To run with the default configuration (Gemma on PyTorch): python3 -m lit_nlp.examples.prompt_debugging.server -- \ - --models=gemma:gemma_1.1_instruct_2b_en \ + --models=gemma:google/gemma-2b-it \ --alsologtostderr MODELS: @@ -19,19 +19,15 @@ strings, and remember the number of models loaded will be limited by the memory available on your accelerator. -| Model | dl_framework | dl_backend=tensorflow Preset | dl_backend=torch Preset | -| ------- | ------------ | ---------------------------- | ------------------------------------ | -| Gemma | kerasnlp | gemma_1.1_instruct_7b_en | gemma_1.1_instruct_7b_en | -| Gemma | transformers | Unavailable | google/gemma-1.1-7b-it | -| Llama 2 | kerasnlp | llama2_instruct_7b_en | llama2_instruct_7b_en | -| Llama 2 | transformers | Unavailable | meta-llama/Llama-2-7b-hf | -| Mistral | kerasnlp | mistral_instruct_7b_en | mistral_instruct_7b_en | -| Mistral | transformers | Unavailable | mistralai/Mistral-7B-Instruct-v0.2 | +| Model | Preset | +| ------- | ---------------------------------- | +| Gemma | google/gemma-1.1-7b-it | +| Llama 2 | meta-llama/Llama-2-7b-hf | +| Mistral | mistralai/Mistral-7B-Instruct-v0.2 | -Additional model presets can be found at the following locations, though +Additional model presets can be found at the following location, though compatibility with the LIT model wrappers is not guaranteed: -* KerasNLP: https://keras.io/api/keras_nlp/models/ * HuggingFace Transformers: https://huggingface.co/models DATASETS: @@ -79,37 +75,14 @@ "The number of examples to process per batch.", ) -_SUPPORTED_FRAMEWORKS = ("kerasnlp", "transformers") -_DL_FRAMEWORK = flags.DEFINE_enum( - "dl_framework", - models.DEFAULT_DL_FRAMEWORK, - _SUPPORTED_FRAMEWORKS, - "The deep learning framework that loads and runs the model on the backend." - " This server will attempt to load all models specified by the --models" - " flag with the configured framework, incompatibilities will result in" - " errors.", -) - -_DL_RUNTIME = flags.DEFINE_enum( - "dl_runtime", - models.DEFAULT_DL_RUNTIME, - # TODO(b/333373960): Add "jax" once JAX salience is supported. - ("tensorflow", "torch"), - "The deep learning backend framework that the model runs on. All models" - " loaded by this server will use the same backend, incompatibilities will" - " result in errors.", -) - _MODELS = flags.DEFINE_list( "models", models.DEFAULT_MODELS, "Models to load, as :. Path can be a URL, a local file path, or" - " the name of a preset for the configured Deep Learning framework (either" - " KerasNLP or HuggingFace Transformers; see --dl_framework for more). This" - " demo is tested with Gemma, GPT2, Llama, and Mistral on all supported" - " --dl_framework values. Other models should work, but adjustments might be" - " needed on their tokenizers (e.g., to define custom pad_token" - " when eos_token is not available to use as pad_token).", + " the name of a HuggingFace Transformers model. This demo is tested with" + " Gemma, GPT2, Llama, and Mistral. Other models should work, but" + " adjustments might be needed on their tokenizers (e.g., to define custom" + " pad_token when eos_token is not available to use as pad_token).", ) _PRECISION = flags.DEFINE_enum( @@ -161,8 +134,6 @@ def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]: lit_demo = dev_server.Server( models=models.get_models( models_config=_MODELS.value, - dl_framework=_DL_FRAMEWORK.value, - dl_runtime=_DL_RUNTIME.value, precision=_PRECISION.value, batch_size=_BATCH_SIZE.value, max_length=_SEQUENCE_LENGTH.value, @@ -172,8 +143,6 @@ def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]: ), layouts=layouts.PROMPT_DEBUGGING_LAYOUTS, model_loaders=models.get_model_loaders( - dl_framework=_DL_FRAMEWORK.value, - dl_runtime=_DL_RUNTIME.value, batch_size=_BATCH_SIZE.value, max_length=_SEQUENCE_LENGTH.value, ), diff --git a/lit_nlp/examples/prompt_debugging/transformers_lms.py b/lit_nlp/examples/prompt_debugging/transformers_lms.py index 7a226691..c904e340 100644 --- a/lit_nlp/examples/prompt_debugging/transformers_lms.py +++ b/lit_nlp/examples/prompt_debugging/transformers_lms.py @@ -7,11 +7,9 @@ hidden states and attention. """ from collections.abc import Sequence -import enum import functools from typing import Any, Mapping -from absl import logging from lit_nlp.api import model as lit_model from lit_nlp.api import types as lit_types from lit_nlp.examples.prompt_debugging import constants as pd_constants @@ -19,46 +17,20 @@ from lit_nlp.lib import file_cache from lit_nlp.lib import utils import numpy as np +import torch import transformers -# pylint: disable=g-import-not-at-top -# pytype: disable=import-error -try: - import tensorflow as tf -except (ModuleNotFoundError, ImportError): - logging.warning("TensorFlow is not available.") - -try: - import torch -except (ModuleNotFoundError, ImportError): - logging.warning("PyTorch is not available.") -# pytype: enable=import-error -# pylint: enable=g-import-not-at-top - _PYTORCH = "torch" -_TENSORFLOW = "tensorflow" -# HuggingFace uses two letter abbreviations for pytorch and tensorflow. +# HuggingFace uses two letter abbreviations for pytorch. _HF_PYTORCH = "pt" -_HF_TENSORFLOW = "tf" - - -@enum.unique -class MLFramework(enum.Enum): - """The supported deep learning frameworks.""" - - PT = _PYTORCH - TF = _TENSORFLOW - -SUPPORTED_ML_RUNTIMES = [framework.value for framework in MLFramework] +SUPPORTED_ML_RUNTIMES = [_PYTORCH] class HFBaseModel(lit_model.BatchedModel): """Base class for HF generative, salience, tokenizer model wrappers.""" - # Enum str values for entries in MLFramework, used for init_spec and logging. - @property def num_layers(self): return self.model.config.n_layer @@ -93,8 +65,7 @@ def __init__( model_name_or_path: gpt2, gpt2-medium, gpt2-large, distilgpt2, meta-llama/Llama-2-7b-hf, mistralai/Mistral-7B-v0.1, etc. batch_size: the number of items to process per `predict_minibatch` call. - framework: the deep learning framework, only "tensorflow" and "torch" - are supported. + framework: the deep learning framework, only "torch" is supported. model: an initialized transformer model. tokenizer: an initialized tokenizer. """ @@ -104,14 +75,10 @@ def __init__( self.model = model self.tokenizer = tokenizer # Check if the HF model object's framework is supported here. - if model.framework == _HF_PYTORCH: - self.framework = MLFramework.PT - elif model.framework == _HF_TENSORFLOW: - self.framework = MLFramework.TF - else: + if model.framework != _HF_PYTORCH: raise ValueError( f"The HuggingFace model framework `{model.framework}` is not" - " supported." + " supported; only PyTorch ('pt') is supported." ) else: # Normally path is a directory; if it's an archive file, download and @@ -136,25 +103,18 @@ def __init__( # AutoTokenizer.from_pretrained() above it will create a new token with # with id = max_vocab_length and cause out-of-bounds errors in # the embedding lookup. - if framework == _PYTORCH: - auto_model = transformers.AutoModelForCausalLM - self.framework = MLFramework.PT - elif framework == _TENSORFLOW: - auto_model = transformers.TFAutoModelForCausalLM - self.framework = MLFramework.TF - else: + if framework != _PYTORCH: raise ValueError( f"The provided value `{framework}` for arg `framework` is not" f" supported, please choose from {SUPPORTED_ML_RUNTIMES}." ) - self.model = auto_model.from_pretrained( + self.model = transformers.AutoModelForCausalLM.from_pretrained( model_name_or_path, output_hidden_states=True, output_attentions=False, ) - if self.framework == MLFramework.PT: - self.device = "cuda:0" if torch.cuda.is_available() else "cpu" - self.model = self.model.to(self.device) + self.device = "cuda:0" if torch.cuda.is_available() else "cpu" + self.model = self.model.to(self.device) self.embedding_table = self.model.get_input_embeddings() self.tokenizer.pad_token = self.tokenizer.eos_token self.batch_size = batch_size @@ -263,22 +223,18 @@ def _get_batched_outputs( Returns: a dict of the model outputs, including the generated texts and auxiliary - data in numpy arrays (could come from torch or tensorflow, depending on - the transformer backend). + data in numpy arrays. """ encoded_inputs = self.tokenizer( [ex["prompt"] for ex in inputs], - return_tensors=( - _HF_PYTORCH if self.framework == MLFramework.PT else _HF_TENSORFLOW - ), + return_tensors=_HF_PYTORCH, add_special_tokens=True, padding="longest", truncation="longest_first", ) batch_size, ntok_in = encoded_inputs["input_ids"].shape - if self.framework == MLFramework.PT: - encoded_inputs = encoded_inputs.to(self.device) + encoded_inputs = encoded_inputs.to(self.device) outputs = self.model.generate(**encoded_inputs, max_length=self.max_length) @@ -291,12 +247,9 @@ def _get_batched_outputs( outputs[:, -ntok_out:], skip_special_tokens=True ) - if self.framework == MLFramework.PT: - with torch.no_grad(): - # Input embeddings: [batch_size, num_tokens, emb_dim] - embeddings = self.embedding_table(outputs).cpu().to(torch.float) - else: - embeddings = self.embedding_table(outputs) + with torch.no_grad(): + # Input embeddings: [batch_size, num_tokens, emb_dim] + embeddings = self.embedding_table(outputs).cpu().to(torch.float) return { "embs": embeddings.numpy(), @@ -351,71 +304,6 @@ def _left_pad_target_masks(self, seq_length, target_masks): ) return padded_target_masks - def _pred_tf(self, encoded_inputs, target_masks): - """Predicts one batch of tokenized text using TF. - - Also performs some batch-level post-processing in TF. - Single-example postprocessing is done in _postprocess(), and operates on - numpy arrays. - - Args: - encoded_inputs: output of self.tokenizer() - target_masks: list(array_like) of binary (0/1) masks for each input - - Returns: - payload: Dictionary with items described above, each as single Tensor. - """ - input_ids = encoded_inputs["input_ids"] - - # [batch_size, num_tokens]; ignore the last one in each row. - target_ids = tf.roll(input_ids, shift=-1, axis=1) - ## - # Process target masks - padded_target_masks = tf.constant( - self._left_pad_target_masks(target_ids.shape[1], target_masks), - dtype=tf.bool, - ) - # Shift masks back so they align with target_ids. - loss_mask = tf.roll(padded_target_masks, shift=-1, axis=1) - - with tf.GradientTape(watch_accessed_variables=False) as tape: - # We need to run the embedding layer ourselves so we can trace it. - # See here for how the model normally does this: - # https://github.com/huggingface/transformers/blob/v4.29.2/src/transformers/models/gpt2/modeling_tf_gpt2.py#L450 - embs = self.embedding_table(input_ids) - tape.watch(embs) - - out = self.model( - input_ids=None, - inputs_embeds=embs, - attention_mask=encoded_inputs["attention_mask"], - ) - - loss_fn = tf.keras.losses.SparseCategoricalCrossentropy( - from_logits=True, reduction="none" - ) - # [batch_size, num_tokens] - per_token_loss = loss_fn(target_ids, out.logits) - masked_loss = per_token_loss * tf.cast(loss_mask, per_token_loss.dtype) - - grads = tape.gradient( - masked_loss, embs - ) # [batch_size, num_tokens, hdim] - - grad_l2 = tf.norm(grads, axis=2) # [batch_size, num_tokens] - grad_dot_input = tf.reduce_sum( - grads * embs, axis=2 - ) # [batch_size, num_tokens] - - batched_outputs = { - "input_ids": input_ids, - "attention_mask": encoded_inputs["attention_mask"], - pd_constants.FieldNames.GRAD_NORM: grad_l2, - pd_constants.FieldNames.GRAD_DOT_INPUT: grad_dot_input, - } - - return batched_outputs - def _pred_pt(self, encoded_inputs, target_masks): """Predicts one batch of tokenized text using PyTorch. @@ -504,9 +392,7 @@ def predict_minibatch(self, inputs): ] encoded_inputs = self.tokenizer( texts, - return_tensors=_HF_PYTORCH - if self.framework == MLFramework.PT - else _HF_TENSORFLOW, + return_tensors=_HF_PYTORCH, add_special_tokens=True, padding="longest", truncation="longest_first", @@ -516,10 +402,7 @@ def predict_minibatch(self, inputs): ] # Get the predictions. - if self.framework == MLFramework.PT: - batched_outputs = self._pred_pt(encoded_inputs, target_masks) - else: - batched_outputs = self._pred_tf(encoded_inputs, target_masks) + batched_outputs = self._pred_pt(encoded_inputs, target_masks) # Convert to numpy for post-processing. detached_outputs = {k: v.numpy() for k, v in batched_outputs.items()} @@ -560,16 +443,14 @@ def predict_minibatch(self, inputs): ] encoded_inputs = self.tokenizer( texts, - return_tensors=_HF_PYTORCH - if self.framework == MLFramework.PT - else _HF_TENSORFLOW, + return_tensors=_HF_PYTORCH, add_special_tokens=True, padding="longest", truncation="longest_first", ) batched_outputs = { - "input_ids": encoded_inputs["input_ids"], - "attention_mask": encoded_inputs["attention_mask"], + "input_ids": encoded_inputs["input_ids"].cpu(), + "attention_mask": encoded_inputs["attention_mask"].cpu(), } # Convert to numpy for post-processing. detached_outputs = {k: v.numpy() for k, v in batched_outputs.items()} diff --git a/lit_nlp/examples/prompt_debugging/transformers_lms_int_test.py b/lit_nlp/examples/prompt_debugging/transformers_lms_int_test.py index 9c7cde6f..296ed7a4 100644 --- a/lit_nlp/examples/prompt_debugging/transformers_lms_int_test.py +++ b/lit_nlp/examples/prompt_debugging/transformers_lms_int_test.py @@ -6,31 +6,27 @@ _MAX_LENGTH = 32 +_GPT2_PT_PATH = ( + "https://storage.googleapis.com/what-if-tool-resources/lit-models/" + "gpt2-pt.tar.gz" +) + def _tokenize_text( - text: str, tokenizer: tokenization_utils.PreTrainedTokenizer, framework: str + text: str, tokenizer: tokenization_utils.PreTrainedTokenizer ) -> tokenization_utils.BatchEncoding: - return_tensors_type = ( - transformers_lms._HF_PYTORCH - if framework == transformers_lms.MLFramework.PT.value - else transformers_lms._HF_TENSORFLOW - ) return tokenizer( text, - return_tensors=return_tensors_type, + return_tensors=transformers_lms._HF_PYTORCH, add_special_tokens=True, ) def _get_text_mean_embeddings( - text: str, model: transformers_lms.HFBaseModel, framework: str + text: str, model: transformers_lms.HFBaseModel ) -> np.ndarray: - tokens = _tokenize_text( - text=text, tokenizer=model.tokenizer, framework=framework - ) - embeddings = model.embedding_table(tokens["input_ids"]) - if framework == transformers_lms.MLFramework.PT.value: - embeddings = embeddings.detach() + tokens = _tokenize_text(text=text, tokenizer=model.tokenizer) + embeddings = model.embedding_table(tokens["input_ids"]).detach() mean_embeddings = np.mean(embeddings.numpy()[0], axis=0) return mean_embeddings @@ -38,22 +34,10 @@ def _get_text_mean_embeddings( class TransformersLMSGeneration(parameterized.TestCase): """Test that model classes can predict.""" - @parameterized.named_parameters( - dict( - testcase_name="tensorflow_framework", - framework=transformers_lms.MLFramework.TF.value, - model_path="https://storage.googleapis.com/what-if-tool-resources/lit-models/gpt2.tar.gz", - ), - dict( - testcase_name="pytorch_framework", - framework=transformers_lms.MLFramework.PT.value, - model_path="https://storage.googleapis.com/what-if-tool-resources/lit-models/gpt2-pt.tar.gz", - ), - ) - def test_gpt2_generation_output(self, framework, model_path): + def test_gpt2_generation_output(self): model = transformers_lms.HFGenerativeModel( - model_name_or_path=model_path, - framework=framework, + model_name_or_path=_GPT2_PT_PATH, + framework=transformers_lms.SUPPORTED_ML_RUNTIMES[0], max_length=_MAX_LENGTH, ) model_in = [{"prompt": "Today is"}, {"prompt": "What is the color of"}] @@ -74,10 +58,10 @@ def test_gpt2_generation_output(self, framework, model_path): ): for cur_input, cur_output in zip(model_in, model_out): expected_input_embeddings = _get_text_mean_embeddings( - text=cur_input["prompt"], model=model, framework=framework + text=cur_input["prompt"], model=model ) expected_output_embeddings = _get_text_mean_embeddings( - text=cur_output["response"], model=model, framework=framework + text=cur_output["response"], model=model ) np.testing.assert_array_almost_equal( expected_input_embeddings, @@ -88,24 +72,12 @@ def test_gpt2_generation_output(self, framework, model_path): cur_output["response_embeddings"], ) - @parameterized.named_parameters( - dict( - testcase_name="tensorflow_framework", - framework=transformers_lms.MLFramework.TF.value, - model_path="https://storage.googleapis.com/what-if-tool-resources/lit-models/gpt2.tar.gz", - ), - dict( - testcase_name="pytorch_framework", - framework=transformers_lms.MLFramework.PT.value, - model_path="https://storage.googleapis.com/what-if-tool-resources/lit-models/gpt2-pt.tar.gz", - ), - ) def test_gpt2_batched_generation_has_correct_input_and_output_token_lengths( - self, framework, model_path + self, ): model = transformers_lms.HFGenerativeModel( - model_name_or_path=model_path, - framework=framework, + model_name_or_path=_GPT2_PT_PATH, + framework=transformers_lms.SUPPORTED_ML_RUNTIMES[0], max_length=_MAX_LENGTH, ) model_in = [{"prompt": "Today is"}, {"prompt": "What is the color of"}] @@ -114,7 +86,6 @@ def test_gpt2_batched_generation_has_correct_input_and_output_token_lengths( _tokenize_text( text=input_dict["prompt"], tokenizer=model.tokenizer, - framework=framework, ) for input_dict in model_in ] diff --git a/lit_nlp/examples/tools/glue_trainer.py b/lit_nlp/examples/tools/glue_trainer.py index eb9d483a..01e47eb3 100644 --- a/lit_nlp/examples/tools/glue_trainer.py +++ b/lit_nlp/examples/tools/glue_trainer.py @@ -14,7 +14,7 @@ Note: you don't have to use this trainer to use LIT; the classifier implementation is just a wrapper around HuggingFace Transformers, using -AutoTokenizer, AutoConfig, and TFAutoModelForSequenceClassification, and can +AutoTokenizer, AutoConfig, and AutoModelForSequenceClassification, and can load anything compatible with those classes. """ @@ -27,9 +27,6 @@ from lit_nlp.examples.glue import data as glue_data from lit_nlp.examples.glue import models as glue_models from lit_nlp.lib import serialize -import tf_keras as keras - -os.environ["TF_USE_LEGACY_KERAS"] = "1" _ENCODER_NAME = flags.DEFINE_string( "encoder_name", "bert-base-uncased", @@ -47,37 +44,6 @@ FLAGS = flags.FLAGS -def history_to_dict(keras_history): - return { - "epochs": keras_history.epoch, - "history": keras_history.history, - "params": keras_history.params, - "optimizer_params": keras_history.model.optimizer.get_config(), - } - - -class EpochSaverCallback(keras.callbacks.Callback): - """Save model at the beginning of training and after every epoch. - - Similar to keras.callbacks.ModelCheckpoint, but this allows us to specify - a custom save fn to call, such as the HuggingFace model.save() which writes - .h5 files and config information. - """ - - def __init__(self, save_path_base: str, save_fn=None): - super().__init__() - self.save_path_base = save_path_base - self.save_fn = save_fn or self.model.save - - def on_train_begin(self, logs=None): - self.on_epoch_end(-1, logs=logs) # write epoch-0 - - def on_epoch_end(self, epoch, logs=None): - # Save path 1-indexed = # of completed epochs. - save_path = os.path.join(self.save_path_base, f"epoch-{epoch+1}") - self.save_fn(save_path) - - def train_and_save(model, train_data, val_data, @@ -85,26 +51,28 @@ def train_and_save(model, save_intermediates=False, **train_kw): """Run training and save model.""" - # Set up logging for TensorBoard. To view, run: - # tensorboard --log_dir=/tensorboard - keras_callbacks = [ - keras.callbacks.TensorBoard( - log_dir=os.path.join(train_path, "tensorboard") - ) - ] - if save_intermediates: - keras_callbacks.append(EpochSaverCallback(train_path, save_fn=model.save)) + + def epoch_saver(epoch, logs): + del logs # unused + if not save_intermediates: + return + # Save path 1-indexed = # of completed epochs. + save_path = os.path.join(train_path, f"epoch-{epoch+1}") + os.makedirs(save_path, exist_ok=True) + model.save(save_path) + history = model.train( train_data.examples, validation_inputs=val_data.examples, - keras_callbacks=keras_callbacks, + on_epoch_end=epoch_saver, **train_kw) # Save training history too, since this is human-readable and more concise # than the TensorBoard log files. + os.makedirs(train_path, exist_ok=True) with open(os.path.join(train_path, "train.history.json"), "w") as fd: # Use LIT's custom JSON encoder to handle dicts containing NumPy data. - fd.write(serialize.to_json(history_to_dict(history), simple=True, indent=2)) + fd.write(serialize.to_json(history, simple=True, indent=2)) model.save(train_path) logging.info("Saved model files: \n %s", diff --git a/lit_nlp/examples/tydi/README.md b/lit_nlp/examples/tydi/README.md index f8b3508d..944fa600 100644 --- a/lit_nlp/examples/tydi/README.md +++ b/lit_nlp/examples/tydi/README.md @@ -3,7 +3,7 @@ TyDi QA Demo for the Learning Interpretability Tool This demo showcases how LIT can be used to a multilingual question-answering model trained on the [TyDi QA dataset](https://doi.org/10.1162/tacl_a_00317) -using FLAX. +using PyTorch. You will need a stand-alone virtual environment for the Python libraries, which you can set up using the following commands from the root of the LIT repo. diff --git a/lit_nlp/examples/tydi/data.py b/lit_nlp/examples/tydi/data.py index 8ad0b912..a36c08d2 100644 --- a/lit_nlp/examples/tydi/data.py +++ b/lit_nlp/examples/tydi/data.py @@ -1,11 +1,10 @@ """Data loaders for Question answering model.""" -import re +import json +from huggingface_hub import hf_hub_download from lit_nlp.api import dataset as lit_dataset from lit_nlp.api import dtypes from lit_nlp.api import types as lit_types -import tensorflow_datasets as tfds - TYDI_LANG_VOCAB = [ 'english', @@ -19,37 +18,89 @@ 'finnish', ] +# TyDi QA gold passages, mirrored per-language on HuggingFace. +# See https://huggingface.co/datasets/khalidalt/tydiqa-goldp +_TYDI_REPO_ID = 'khalidalt/tydiqa-goldp' + +_LANG_BY_CODE = { + 'ar': 'arabic', + 'bn': 'bengali', + 'en': 'english', + 'fi': 'finnish', + 'id': 'indonesian', + 'ja': 'japanese', + 'ko': 'korean', + 'ru': 'russian', + 'sw': 'swahili', + 'te': 'telugu', + 'th': 'thai', +} + + +def _load_jsonl(split_dir: str, language: str) -> list[dict]: + path = hf_hub_download( + repo_id=_TYDI_REPO_ID, + repo_type='dataset', + filename=f'{split_dir}/{language}-{split_dir}.jsonl', + ) + with open(path, encoding='utf-8') as f: + return [json.loads(line) for line in f] + + +def _char_offset(text: str, byte_offset: int) -> int: + """Convert a UTF-8 byte offset in text to a character offset.""" + return len(text.encode('utf-8')[:byte_offset].decode('utf-8', 'ignore')) + class TyDiQA(lit_dataset.Dataset): """TyDiQA dataset.""" def __init__(self, split: str, max_examples=-1): - ds = tfds.load('tydi_qa', split=split) + if split == 'train': + split_dir = 'train' + elif split == 'validation': + split_dir = 'dev' + elif split.startswith('validation-'): + code = split[len('validation-') :] + if code not in _LANG_BY_CODE: + raise ValueError(f'Unknown language code in split: {split}') + split_dir = 'dev' + else: + raise ValueError(f'Unsupported split: {split}') + + languages = ( + [_LANG_BY_CODE[split[len('validation-') :]]] + if split.startswith('validation-') + else sorted(_LANG_BY_CODE.values()) + ) # populate this with data records self._examples = [] - for row in ds.take(max_examples): - answers_text = row['answers']['text'].numpy() - answers_start = [row['answers']['answer_start'].numpy()[0]] - answers = [] - # gets language id example: finnish--9069599​462862564793-0 - language_id = row['id'].numpy().decode('utf-8') - alpha_chars_filter = re.findall(r'[a-z]', language_id) - language = ''.join(str(r) for r in alpha_chars_filter) - - for label, start in zip(answers_text, answers_start): - span = dtypes.SpanLabel(start, start + len(label), align='context') - answers.append( - dtypes.AnnotationCluster(label=label.decode('utf-8'), spans=[span]) - ) + for language in languages: + for row in _load_jsonl(split_dir, language): + context = row['passage_text'] + answers = [] + for answer in row['answers']: + label = answer['text'] + start = _char_offset(context, answer['start_byte']) + span = dtypes.SpanLabel(start, start + len(label), align='context') + answers.append( + dtypes.AnnotationCluster(label=label, spans=[span]) + ) - self._examples.append({ - 'answers_text': answers, - 'title': row['title'].numpy().decode('utf-8'), - 'context': row['context'].numpy().decode('utf-8'), - 'question': row['question'].numpy().decode('utf-8'), - 'language': language, - }) + self._examples.append({ + 'answers_text': answers, + 'title': row['document_title'], + 'context': context, + 'question': row['question_text'], + 'language': row['language'], + }) + if 0 <= max_examples <= len(self._examples): + break + if 0 <= max_examples <= len(self._examples): + break + if max_examples >= 0: + self._examples = self._examples[:max_examples] def spec(self) -> lit_types.Spec: return { diff --git a/lit_nlp/examples/tydi/model.py b/lit_nlp/examples/tydi/model.py index 55001377..afc6fad0 100644 --- a/lit_nlp/examples/tydi/model.py +++ b/lit_nlp/examples/tydi/model.py @@ -4,16 +4,16 @@ from lit_nlp.api import types as lit_types from lit_nlp.examples.tydi import data as tydi_data import numpy as np +import torch import transformers - -_BertTokenizer = transformers.BertTokenizer -_FlaxBertForQuestionAnswering = transformers.FlaxBertForQuestionAnswering +_AutoTokenizer = transformers.AutoTokenizer +_AutoModelForQuestionAnswering = transformers.AutoModelForQuestionAnswering _JsonDict = lit_types.JsonDict class TyDiModel(lit_model.Model): - """Question Answering Jax model based on TyDiQA Dataset.""" + """Question Answering PyTorch model based on TyDiQA Dataset.""" def __init__( self, @@ -23,10 +23,13 @@ def __init__( **unused_kw, ): super().__init__() - self.tokenizer = tokenizer or _BertTokenizer.from_pretrained(model_name) - self.model = model or _FlaxBertForQuestionAnswering.from_pretrained( + self.tokenizer = tokenizer or _AutoTokenizer.from_pretrained(model_name) + self.model = model or _AutoModelForQuestionAnswering.from_pretrained( model_name ) + self.device = "cuda:0" if torch.cuda.is_available() else "cpu" + self.model.to(self.device) + self.model.eval() def _segment_slicers(self, tokens: list[str]): """Slicers along the tokens dimension for each segment. @@ -38,7 +41,7 @@ def _segment_slicers(self, tokens: list[str]): tokens: [num_tokens], including special tokens Returns: - (slicer_a, slicer_b), slice objects + (slicer_question, slicer_context), slice objects """ try: split_point = tokens.index(self.tokenizer.sep_token) @@ -59,19 +62,20 @@ def predict(self, inputs: Iterable[_JsonDict], **kw) -> Iterable[_JsonDict]: for inp in inputs: tokenized_text = self.tokenizer( - inp["question"], inp["context"], return_tensors="jax", padding=True - ) - results = self.model( - **tokenized_text, output_hidden_states=True - ) - answer_start_index = results.start_logits.argmax() - answer_end_index = results.end_logits.argmax() + inp["question"], inp["context"], return_tensors="pt", padding=True + ).to(self.device) + with torch.no_grad(): + results = self.model( + **tokenized_text, output_hidden_states=True + ) + answer_start_index = results.start_logits.argmax().item() + answer_end_index = results.end_logits.argmax().item() predict_answer_tokens = tokenized_text.input_ids[ 0, answer_start_index : answer_end_index + 1 ] # get id's for question & context - tokens = np.asarray(tokenized_text["input_ids"]) + tokens = np.asarray(tokenized_text["input_ids"].cpu()) # convert id's to tokens total_tokens = self.tokenizer.convert_ids_to_tokens(tokens[0]) # split by question & context diff --git a/lit_nlp/examples/tydi/requirements.txt b/lit_nlp/examples/tydi/requirements.txt index 41c3762c..53fb3409 100644 --- a/lit_nlp/examples/tydi/requirements.txt +++ b/lit_nlp/examples/tydi/requirements.txt @@ -13,11 +13,4 @@ # limitations under the License. # ============================================================================== --r ../../../requirements.txt - -flax>=0.5.3 -jax==0.4.6 -jaxlib==0.4.6 -tensorflow>=2.10.0,<2.16.0 -tensorflow-datasets>=4.9.0 -transformers>=4.27.1 +-r ../../../requirements_examples_discriminative_ai.txt diff --git a/lit_nlp/lib/image_utils.py b/lit_nlp/lib/image_utils.py index 579752f3..d185dd21 100644 --- a/lit_nlp/lib/image_utils.py +++ b/lit_nlp/lib/image_utils.py @@ -3,7 +3,7 @@ import io from typing import Optional -import matplotlib.cm as plt_cm +from matplotlib import colormaps as plt_cm import numpy as np from PIL import Image as PILImage from PIL import ImageEnhance as PILImageEnhance @@ -125,7 +125,7 @@ def overlay_pixel_saliency(image_str: str, saliency: np.ndarray, cm_name: str, norm_saliency = normalize_unsigned_saliency(clipped_saliency) # Map saliency to RGB values. - cm = plt_cm.get_cmap(cm_name) + cm = plt_cm[cm_name] saliency_bytes = cm(norm_saliency) # Assign alpha values. @@ -134,7 +134,7 @@ def overlay_pixel_saliency(image_str: str, saliency: np.ndarray, cm_name: str, alphas = map(lambda e: abs(e - 0.5) * 2, norm_saliency.flatten()) else: alphas = map(lambda e: 1.0 - e, norm_saliency.flatten()) - alphas = np.reshape(list(alphas), newshape=norm_saliency.shape) + alphas = np.reshape(list(alphas), norm_saliency.shape) else: alphas = 1.0 diff --git a/lit_nlp/package.json b/lit_nlp/package.json index 42a26bdd..49008f64 100644 --- a/lit_nlp/package.json +++ b/lit_nlp/package.json @@ -1,6 +1,6 @@ { "name": "lit-client", - "version": "1.3.1", + "version": "1.4.0", "description": "Client app for LIT, the Learning Interpretability Tool", "engines": { "node": ">=18.0.0" diff --git a/pyproject.toml b/pyproject.toml index 2a27db38..8d6a5599 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "lit-nlp" # LINT.IfChange -version = "1.3.1" +version = "1.4.0" # LINT.ThenChange(./lit_nlp/package.json) authors = [ { name="Google, LLC", email="lit-dev@google.com" } @@ -13,7 +13,7 @@ authors = [ description = "🔥LIT: The Learning Interpretability Tool" readme = "README.md" license = { file = "LICENSE" } -requires-python = ">=3.9" +requires-python = ">=3.12,<4" # LINT.IfChange dependencies = [ "absl-py>=1.4.0", @@ -24,10 +24,10 @@ dependencies = [ "google-cloud-translate>=3.11.1", "ipython>=7.34.0", "Levenshtein>=0.21.1", - "matplotlib>=3.7.1", + "matplotlib>=3.8.0", "ml-collections>=0.1.1", - "numpy>=1.24.1,<2.0.0", - "pandas[output-formatting]>=2.0.3", + "numpy>=2.1.0,<3", + "pandas[output-formatting]>=2.2.0,<4", "Pillow>=10.0.0", "portpicker>=1.5.2", "requests>=2.31.0", @@ -36,7 +36,7 @@ dependencies = [ "saliency>=0.1.3", "scikit-learn>=1.0.2", "scipy>=1.10.1", - "shap>=0.42.0,<0.46.0", + "shap>=0.52.0", "six>=1.16.0", "termcolor>=2.3.0", "tqdm>=4.64.0", @@ -53,7 +53,9 @@ classifiers = [ "Operating System :: OS Independent", "Programming Language :: JavaScript", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Scientific/Engineering :: Information Analysis", "Topic :: Scientific/Engineering :: Visualization", @@ -72,48 +74,44 @@ keywords = [ "deep learning", "visualization", "visual analytics", - "tensorflow", - "tf", "torch", ] [project.optional-dependencies] +# NOTE: Extras are ADDITIVE and disjoint: no package is listed in more than +# one extra, and extras never depend on each other (poetry#10447 rejects +# self-referential extras, and duplicated deps across extras trigger a +# full re-resolution per overlap in Poetry's solver). Install the combination +# you need, e.g. 'lit-nlp[examples-common,examples-discriminative-ai]'. +# The "-r" chains in requirements_examples_*.txt express the intended combos; +# keep the union of all extras in sync with the union of those files. # LINT.IfChange examples-common = [ "gunicorn>=20.1.0", - "tensorflow>=2.16.0", - "transformers>=4.27.1", + "torch>=2.0.0", + "transformers>=4.27.1,<5", ] # LINT.ThenChange(./requirements_examples_common.txt) # LINT.IfChange examples-discriminative-ai = [ - "lit-nlp[examples-common]", - "tensorflow-datasets>=4.9.0", - "tf-keras>=2.16", + "datasets>=3.0,<4", + "huggingface-hub>=0.25.0,<1.0", ] # LINT.ThenChange(./requirements_examples_discriminative_ai.txt) # LINT.IfChange examples-generative-ai = [ - "lit-nlp[examples-common]", "google-cloud-aiplatform>=1.60.0", - "keras>=3.0.0", - "keras-nlp>=0.14.0", - "sentencepiece==0.1.99", - "tensorflow-text>=2.16.0", - "torch>=2.0.0", + "sentencepiece>=0.2.0,<0.3", "vertexai>=1.49.0", ] # LINT.ThenChange(./requirements_examples_generative_ai.txt) umap = [ - "umap-learn==0.5.6" + "umap-learn>=0.5.9,<0.6" ] # LINT.IfChange test = [ - "lit-nlp[examples-discriminative-ai]", - "lit-nlp[examples-generative-ai]", - "lit-nlp[umap]", "lime==0.2.0.1", - "pytest>=7.4.0,<8.0.0", + "pytest>=8.0,<10", "webtest>=2.0", ] # LINT.ThenChange(./requirements_test.txt) diff --git a/requirements.txt b/requirements.txt index 9c4707b0..7b90f9fd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -21,10 +21,10 @@ filelock>=3.12.3 google-cloud-translate>=3.11.1 ipython>=7.34.0 Levenshtein>=0.21.1 -matplotlib>=3.6.0,<3.9.0 +matplotlib>=3.8.0 ml-collections>=0.1.1 -numpy>=1.24.1,<2.0.0 -pandas[output-formatting]>=2.0.3 +numpy>=2.1.0,<3 +pandas[output-formatting]>=2.2.0,<4 Pillow>=10.0.0 portpicker>=1.5.2 requests>=2.31.0 @@ -33,7 +33,7 @@ sacrebleu>=2.3.1 saliency>=0.1.3 scikit-learn>=1.0.2 scipy>=1.10.1 -shap>=0.42.0,<0.46.0 +shap>=0.52.0 six>=1.16.0 termcolor>=2.3.0 tqdm>=4.64.0 diff --git a/requirements_examples_common.txt b/requirements_examples_common.txt index 511505c9..67fbec12 100644 --- a/requirements_examples_common.txt +++ b/requirements_examples_common.txt @@ -17,6 +17,6 @@ -r requirements.txt gunicorn>=20.1.0 -tensorflow>=2.16.0 -transformers>=4.27.1 +torch>=2.0.0 +transformers>=4.27.1,<5 # LINT.ThenChange(./pyproject.toml) diff --git a/requirements_examples_discriminative_ai.txt b/requirements_examples_discriminative_ai.txt index eba150dd..55ab038b 100644 --- a/requirements_examples_discriminative_ai.txt +++ b/requirements_examples_discriminative_ai.txt @@ -15,6 +15,6 @@ # LINT.IfChange -r requirements_examples_common.txt -tensorflow-datasets>=4.9.0 -tf-keras>=2.16 +datasets>=3.0,<4 +huggingface-hub>=0.25.0,<1.0 # LINT.ThenChange(./pyproject.toml) diff --git a/requirements_examples_generative_ai.txt b/requirements_examples_generative_ai.txt index bb2d5066..a94ef57d 100644 --- a/requirements_examples_generative_ai.txt +++ b/requirements_examples_generative_ai.txt @@ -16,10 +16,6 @@ -r requirements_examples_common.txt google-cloud-aiplatform>=1.60.0 -keras>=3.0.0 -keras-nlp>=0.14.0 -sentencepiece==0.1.99 -tensorflow-text>=2.16.0 -torch>=2.0.0 +sentencepiece>=0.2.0,<0.3 vertexai>=1.49.0 # LINT.ThenChange(./pyproject.toml) diff --git a/requirements_test.txt b/requirements_test.txt index 9e2e0556..4d69ac26 100644 --- a/requirements_test.txt +++ b/requirements_test.txt @@ -17,7 +17,7 @@ -r requirements_examples_generative_ai.txt lime==0.2.0.1 -pytest>=7.4.0,<8.0.0 -umap-learn==0.5.6 +pytest>=8.0,<10 +umap-learn>=0.5.9,<0.6 webtest>=2.0 # LINT.ThenChange(./pyproject.toml) diff --git a/website/sphinx_src/api.md b/website/sphinx_src/api.md index 97ea5b64..4fac4aba 100644 --- a/website/sphinx_src/api.md +++ b/website/sphinx_src/api.md @@ -149,7 +149,7 @@ how to interpret them: open set of values. This implementation uses Pandas to read a TSV file, but you can also use -services like [TensorFlow Datasets](https://www.tensorflow.org/datasets) - +services like [HuggingFace datasets](https://huggingface.co/docs/datasets) - simply wrap them in your `__init__()` function. Note that you can freely add additional features - such as `genre` in the @@ -160,9 +160,9 @@ interpretation components such as custom metrics. ### Transformations The `Dataset` class also supports a limited set of transformations, similar to -TensorFlow's -[tf.data.Dataset](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) -but more limited in scope and aimed at supporting quick iteration: +(other but more limited in scope than) pipelines like +[HuggingFace datasets](https://huggingface.co/docs/datasets) and aimed at +supporting quick iteration: * `Dataset.slice[start:step:end]` will return a new `Dataset` with the same spec and a slice of the datapoints. diff --git a/website/sphinx_src/components.md b/website/sphinx_src/components.md index 90b99dbd..6fd6f251 100644 --- a/website/sphinx_src/components.md +++ b/website/sphinx_src/components.md @@ -23,7 +23,7 @@ out-of-the-box support for a few modeling frameworks, described below. Many of the [open-source LIT examples](https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/) use -HuggingFace Transformers via their TF2/Keras model classes. These give easy +HuggingFace Transformers via their PyTorch model classes. These give easy access to model internals such as embeddings, attention, and gradients, and the LIT wrappers for these support many interpretability methods - such as [integrated gradients](https://arxiv.org/abs/1703.01365) out-of-the-box. @@ -31,23 +31,6 @@ LIT wrappers for these support many interpretability methods - such as These models are a great place to start for small-scale experiments or for working on academic projects. -### TF1.x Estimator - -LIT supports Estimator and other TF1.x models, but the model wrappers can be -more involved due to the need to explicitly manage the graph and sessions. (In -particular: `Estimator.predict()` cannot be used because it reloads the model on -every invocation.) Generally, you'll need to: - -* In your model's `__init__()`, build the graph, create a persistent TF - session, and load the model weights. -* In your `predict()` function, build a feed dict and call `session.run` - directly. - -Alternatively, you can export to a `SavedModel` and load this in an eager mode -runtime. This leads to much simpler code, but may require changes to your -`SavedModel` exporter in order to access model internals like embeddings, -gradients, or attention. - ### Remote or hosted models LIT can easily interact with models hosted via an RPC or HTTP endpoint, @@ -74,7 +57,7 @@ as LIME which involve querying the model on new examples. LIT can load data from almost any format, including TFRecord, Capacitor, SSTable, or even SQL queries, via a custom Python class that implements the -[Dataset API](./api.md#datasets). Many of our demos use TFDS, and the LIT loader +[Dataset API](./api.md#datasets). Many of our demos use [HuggingFace datasets](https://huggingface.co/docs/datasets), and the LIT loader code is only a small wrapper class which performs minimal post-processing. See the [demos page](./demos.md) for specific examples. @@ -462,14 +445,12 @@ If you find this useful in your work, please cite Sequence Salience as: **Code:** Currently, this works out-of-the-box with Gemma, Llama 2, Mistral, and GPT-2, -using either KerasNLP or Transformers. +using HuggingFace Transformers. * LIT-for-Gemma Colab: [`lit_gemma.ipynb`](https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/lit_gemma.ipynb) * Demo binary: https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/prompt_debugging/server.py -* KerasNLP model wrappers: - https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/prompt_debugging/keras_lms.py * Transformers model wrappers: https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/prompt_debugging/transformers_lms.py diff --git a/website/sphinx_src/demos.md b/website/sphinx_src/demos.md index 448c0d23..898b2871 100644 --- a/website/sphinx_src/demos.md +++ b/website/sphinx_src/demos.md @@ -29,7 +29,7 @@ https://pair-code.github.io/lit/demos/. * STS-B textual similarity task (see [Regression / Scoring](#regression-scoring) below). * Switch tasks using the Settings (⚙️) menu. -* BERT models of different sizes, built on HuggingFace TF2 (Keras). +* BERT models of different sizes, built on HuggingFace Transformers (PyTorch). * Supports the widest range of LIT interpretability features: * Model output probabilities, custom thresholds, and multiclass metrics. * Jitter plot of output scores, to find confident examples or ones near @@ -53,7 +53,7 @@ https://pair-code.github.io/lit/tutorials/sentiment * STS-B textual similarity task, predicting scores on a range from 0 (unrelated) to 5 (very similar). -* BERT models built on HuggingFace TF2 (Keras). +* BERT models built on HuggingFace Transformers (PyTorch). * Supports a wide range of LIT interpretability features: * Model output scores and metrics. * Scatter plot of scores and error, and jitter plot of true labels for @@ -70,8 +70,8 @@ https://pair-code.github.io/lit/tutorials/sentiment **Code:** [examples/prompt_debugging/server.py](https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/prompt_debugging/server.py) -* Supports Gemma 2B and 7B models using KerasNLP (with TensorFlow or PyTorch) - and Transformers (with PyTorch). +* Supports Gemma, Llama, Mistral, and other LLMs using HuggingFace + Transformers (with PyTorch). * Interactively debug LLM prompts using [sequence salience](./components.md#sequence-salience). * Multiple salience methods (grad-l2 and grad-dot-input), at multiple @@ -91,7 +91,7 @@ Generative AI Toolkit. **Code:** [examples/penguin/demo.py](https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/penguin/demo.py) * Binary classification on - [penguin dataset](https://www.tensorflow.org/datasets/catalog/penguins). + [penguin dataset](https://allisonhorst.github.io/palmerpenguins/). * Showing using of LIT on non-text data (numeric and categorical features). * Use partial-dependence plots to understand feature importance on individual examples, selections, or the entire evaluation dataset. diff --git a/website/sphinx_src/faq.md b/website/sphinx_src/faq.md index 8e334b24..e4689577 100644 --- a/website/sphinx_src/faq.md +++ b/website/sphinx_src/faq.md @@ -20,7 +20,7 @@ images and tabular data. For examples, see: * [Tabular demo](https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/penguin/demo.py) - multi-class classification on tabular (numeric and categorical string) data, using the - [Palmer Penguins](https://www.tensorflow.org/datasets/catalog/penguins) + [Palmer Penguins](https://allisonhorst.github.io/palmerpenguins/) dataset. For more details, see @@ -168,7 +168,7 @@ LIT is primarily an evaluation/inference-time tool, so we don't provide any official training APIs. However, to facilitate code reuse you can easily add training methods to your model class. In fact, several of our demos do exactly this, using LIT's `Dataset` objects to manage training data along with standard -training APIs (such as Keras' `model.fit()`). See +training APIs (such as PyTorch optimizers and loops). See [`glue/models.py`](https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/glue/models.py) for examples. diff --git a/website/src/index.md b/website/src/index.md index 690c463a..46ced8ad 100644 --- a/website/src/index.md +++ b/website/src/index.md @@ -78,11 +78,11 @@ And more... " %} {% include partials/one-of-three-column title:"Framework agnostic", text: " -TensorFlow 1.x +PyTorch -TensorFlow 2.x +HuggingFace Transformers -PyTorch +scikit-learn Notebook compatibility @@ -104,7 +104,7 @@ And more... {% include partials/home-card image: '/assets/images/LIT_Updates.png', action: 'UPDATES', title: 'Version 1.2', - desc: 'Input salience for text-to-text LLMs, with wrappers for HuggingFace Transformers and KerasNLP models.', + desc: 'Input salience for text-to-text LLMs, with wrappers for HuggingFace Transformers models.', cta-text:"See release notes", link: 'https://github.com/PAIR-code/lit/blob/main/RELEASE.md' external:"true" %} diff --git a/website/src/tutorials/sequence-salience.md b/website/src/tutorials/sequence-salience.md index 3f5ad91b..0a83f8b5 100644 --- a/website/src/tutorials/sequence-salience.md +++ b/website/src/tutorials/sequence-salience.md @@ -52,7 +52,7 @@ All examples in this tutorial use the [Gemma][gemma] LLM as the analysis target. Most of the time, this is Gemma Instruct 2B, but we also use Gemma Instruct 7B in Case Study 3; [more info about variants][gemma_variants] is available online. LIT supports additional LLMs, including [Llama 2][llama] and [Mistral][mistral], -via the HuggingFace Transformers and KerasNLP libraries. +via the HuggingFace Transformers library. This tutorial was adapted from and expands upon LIT's contributions to the [Responsible Generative AI Tookit][rai_toolkit] and the related @@ -421,8 +421,7 @@ for side-by-side comparison and datapoint editing. Salience methods for LLMs is an [active][salience_research_1] [research][salience_research_1] area. The LIT team has provided reference implementations for computing gradient-based salience— -[Grad L2 Norm][grad_norm] and [Grad · Input][grad_dot]—for LLMs in two -popular frameworks: [KerasNLP][lit_keras] and +[Grad L2 Norm][grad_norm] and [Grad · Input][grad_dot]—for LLMs using [HuggingFace Transformers][lit_hf]. There is considerable opportunity to research how the model analysis foundations @@ -472,7 +471,6 @@ helpful guides that can help you develop better prompts, including: [lit_colab]: https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/lit_gemma.ipynb [lit_hf]: https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/prompt_debugging/transformers_lms.py [lit_issues]: https://github.com/PAIR-code/lit/issues -[lit_keras]: https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/prompt_debugging/keras_lms.py [lit_sxs]: ../../documentation/ui_guide.html#comparing-datapoints [llama]: https://llama.meta.com/ [main_toolbar]: ../../documentation/ui_guide.html#main-toolbar