Deprecate jax and tensorflow, remove unused files

Author: RobinPicard

outlines/models/transformers.py

@@ -244,11 +244,29 @@ def __init__(
             and isinstance(model, FlaxPreTrainedModel)
         ):
             self.tensor_library_name = "jax"
+            warnings.warn("""
+                Support for `jax` has been deprecated and will be removed in
+                version 1.4.0 of Outlines. Please use `torch` instead.
+                Transformers models using `jax` do not support structured
+                generation.
+                """,
+                DeprecationWarning,
+                stacklevel=2,
+            )
         elif (
             TFPreTrainedModel is not None
             and isinstance(model, TFPreTrainedModel)
         ):
             self.tensor_library_name = "tensorflow"
+            warnings.warn("""
+                Support for `tensorflow` has been deprecated and will be removed in
+                version 1.4.0 of Outlines. Please use `torch` instead.
+                Transformers models using `tensorflow` do not support structured
+                generation.
+                """,
+                DeprecationWarning,
+                stacklevel=2,
+            )
         else:
             self.tensor_library_name = "torch"
 

outlines/processors/base_logits_processor.py

@@ -28,8 +28,8 @@ def __init__(self, tensor_library_name: str):
         ----------
         tensor_library_name
             The name of the library to use to manipulate tensors. Possible
-            values are "jax", "mlx", "numpy", "tensorflow" and "torch". You
-            must choose the library that your model is using.
+            values are "mlx", "numpy" and "torch". You must choose the library
+            that your model is using.
         """
         # Temporary fix as torch raises a warning that can cause can an error
         # with python 3.12.
@@ -52,7 +52,7 @@ def reset(self):
         needs to be reset for a new generation.
 
         """
-        pass
+        pass # pragma: no cover
 
     @abstractmethod
     def process_logits(

outlines/processors/tensor_adapters/__init__.py

@@ -2,25 +2,19 @@
 
 from typing import Union
 
-from .jax import JAXTensorAdapter
 from .mlx import MLXTensorAdapter
 from .numpy import NumpyTensorAdapter
-from .tensorflow import TensorFlowTensorAdapter
 from .torch import TorchTensorAdapter
 
 
 tensor_adapters = {
-    "jax": JAXTensorAdapter,
     "mlx": MLXTensorAdapter,
     "numpy": NumpyTensorAdapter,
-    "tensorflow": TensorFlowTensorAdapter,
     "torch": TorchTensorAdapter,
 }
 
 TensorAdapterImplementation = Union[
-    JAXTensorAdapter,
     MLXTensorAdapter,
     NumpyTensorAdapter,
-    TensorFlowTensorAdapter,
     TorchTensorAdapter,
 ]

outlines/processors/tensor_adapters/jax.py

@@ -40,25 +40,27 @@ def test_transformers_instantiate_simple():
 
 
 def test_transformers_instantiate_flax_model():
-    model = outlines.from_transformers(
-        transformers.FlaxAutoModelForCausalLM.from_pretrained(TEST_MODEL),
-        transformers.AutoTokenizer.from_pretrained(TEST_MODEL),
-    )
-    assert isinstance(model, Transformers)
-    assert isinstance(model.tokenizer, TransformerTokenizer)
-    assert isinstance(model.type_adapter, TransformersTypeAdapter)
-    assert model.tensor_library_name == "jax"
+    with pytest.warns(DeprecationWarning):
+        model = outlines.from_transformers(
+            transformers.FlaxAutoModelForCausalLM.from_pretrained(TEST_MODEL),
+            transformers.AutoTokenizer.from_pretrained(TEST_MODEL),
+        )
+        assert isinstance(model, Transformers)
+        assert isinstance(model.tokenizer, TransformerTokenizer)
+        assert isinstance(model.type_adapter, TransformersTypeAdapter)
+        assert model.tensor_library_name == "jax"
 
 
 def test_transformers_instantiate_tensorflow_model():
-    model = outlines.from_transformers(
-        transformers.TFAutoModelForCausalLM.from_pretrained(TEST_MODEL),
-        transformers.AutoTokenizer.from_pretrained(TEST_MODEL),
-    )
-    assert isinstance(model, Transformers)
-    assert isinstance(model.tokenizer, TransformerTokenizer)
-    assert isinstance(model.type_adapter, TransformersTypeAdapter)
-    assert model.tensor_library_name == "tensorflow"
+    with pytest.warns(DeprecationWarning):
+        model = outlines.from_transformers(
+            transformers.TFAutoModelForCausalLM.from_pretrained(TEST_MODEL),
+            transformers.AutoTokenizer.from_pretrained(TEST_MODEL),
+        )
+        assert isinstance(model, Transformers)
+        assert isinstance(model.tokenizer, TransformerTokenizer)
+        assert isinstance(model.type_adapter, TransformersTypeAdapter)
+        assert model.tensor_library_name == "tensorflow"
 
 
 def test_transformers_instantiate_mamba():

outlines/processors/tensor_adapters/tensorflow.py

@@ -1,6 +1,5 @@
 from typing import List
 
-import jax.numpy as jnp
 import numpy as np
 import pytest
 import torch
@@ -14,7 +13,7 @@
     HAS_MLX = False
 
 
-libraries = ["numpy", "torch", "jax"]
+libraries = ["numpy", "torch"]
 if HAS_MLX:
     libraries.append("mlx")
 
@@ -43,14 +42,6 @@
         (torch.tensor([1, 2], dtype=torch.float32), torch.tensor([[1, 2], [3, 4]], dtype=torch.int32), AssertionError),
         (torch.tensor([[[1, 2]]], dtype=torch.float32), torch.tensor([[[1, 2]]], dtype=torch.int32), ValueError),
     ],
-    "jax": [
-        (jnp.array([1, 2], dtype=jnp.float32), jnp.array([1, 2], dtype=jnp.int32), None),
-        (jnp.array([[1, 2], [3, 4]], dtype=jnp.float32), jnp.array([[1, 2], [3, 4]], dtype=jnp.int32), None),
-        (jnp.array([1, 2], dtype=jnp.float32), jnp.array([[1, 2]], dtype=jnp.int32), None),
-        (jnp.array([[1, 2]], dtype=jnp.float32), jnp.array([1, 2], dtype=jnp.int32), AssertionError),
-        (jnp.array([1, 2], dtype=jnp.float32), jnp.array([[1, 2], [3, 4]], dtype=jnp.int32), AssertionError),
-        (jnp.array([[[1, 2]]], dtype=jnp.float32), jnp.array([[[1, 2]]], dtype=jnp.int32), ValueError),
-    ],
 }
 if HAS_MLX:
     arrays["mlx"] = [
@@ -76,6 +67,7 @@ def test_base_logits_processor_init(library):
     assert processor.tensor_adapter is not None
     with pytest.raises(NotImplementedError):
         processor = MockLogitsProcessor("foo")
+        processor.reset()
 
 
 @pytest.mark.parametrize("library", libraries)

tests/models/test_transformers.py

@@ -1,18 +1,13 @@
 import pytest
 from pytest import mark
 
-import jax
-import jax.numpy as jnp
 import numpy as np
-import tensorflow as tf
 import torch
 
 from outlines.processors.tensor_adapters import (
     NumpyTensorAdapter,
     TorchTensorAdapter,
     MLXTensorAdapter,
-    JAXTensorAdapter,
-    TensorFlowTensorAdapter,
 )
 
 try:
@@ -25,26 +20,19 @@
 
 
 adapters = {
-    "jax": JAXTensorAdapter(),
     "numpy": NumpyTensorAdapter(),
-    "tensorflow": TensorFlowTensorAdapter(),
     "torch": TorchTensorAdapter(),
 }
 if HAS_MLX:
     adapters["mlx"] = MLXTensorAdapter()
 
-frameworks = ["jax", "numpy", "tensorflow", "torch", "mlx"]
+frameworks = ["numpy", "torch", "mlx"]
 
 def create_tensor(framework, shape, dtype=None):
     if framework == "torch":
         return torch.randn(*shape)
     elif framework == "numpy":
         return np.random.randn(*shape)
-    elif framework == "jax":
-        key = jax.random.PRNGKey(0)
-        return jax.random.normal(key, shape=shape)
-    elif framework == "tensorflow":
-        return tf.random.normal(shape)
     elif framework == "mlx":
         if not HAS_MLX:
             pytest.skip("MLX not available")
@@ -55,10 +43,6 @@ def compare_tensors(framework, tensor1, tensor2):
         return torch.allclose(tensor1, tensor2)
     elif framework == "numpy":
         return np.array_equal(tensor1, tensor2)
-    elif framework == "jax":
-        return jax.numpy.array_equal(tensor1, tensor2)
-    elif framework == "tensorflow":
-        return tf.reduce_all(tf.equal(tensor1, tensor2))
     elif framework == "mlx":
         if not HAS_MLX:
             pytest.skip("MLX not available")
@@ -243,11 +227,6 @@ def test_tensor_adapter_apply_mask(framework):
         mask = torch.randn(2, 3) > 0
     elif framework == "numpy":
         mask = np.random.randn(2, 3) > 0
-    elif framework == "jax":
-        key = jax.random.PRNGKey(0)
-        mask = jax.random.normal(key, shape=(2, 3)) > 0
-    elif framework == "tensorflow":
-        mask = tf.random.normal((2, 3)) > 0
     elif framework == "mlx":
         if not HAS_MLX:
             pytest.skip("MLX not available")