Migrate from TensorFlow to PyTorch; modernize packaging for Python 3.…

.github/workflows/ci.yml

@@ -15,7 +15,7 @@
 
 # Config file for GitHub Actions for CI.
 
 name: CI
 
 # Run on pushes and pull requests on main and dev branches.
 on:
@@ -28,46 +28,46 @@
   workflow_dispatch: null
 
 jobs:
   build-and-test:
     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
       run: pushd lit_nlp && yarn && yarn build && popd
   # TODO(b/235626637): Add front-end unit and e2e tests
 
   docker-build-and-test:
     name: Build docker image
     needs: build-and-test
     runs-on: ubuntu-latest
     steps:
     - name: Set up QEMU
       uses: docker/setup-qemu-action@v2
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v2
     - name: Build Docker image
       uses: docker/build-push-action@v4
       with:
         target: lit-nlp-prod
         tags: lit-nlp:ci-${{ github.sha }}

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"]

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

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).

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.

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))
 

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]))
 

lit_nlp/components/hotflip_int_test.py

@@ -14,22 +14,20 @@
 # ==============================================================================
 """Tests for lit_nlp.components.hotflip."""
 
+import os
+
 from absl.testing import absltest
 from absl.testing import parameterized
 from lit_nlp.components import hotflip
 # TODO(lit-dev): Move glue_models out of lit_nlp/examples
 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'],

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