Merge pull request #764 from shadeMe/chore/merge-master-into-v9

Merge `master` into `v9`

Author: svlandeg

README.md

@@ -2,7 +2,7 @@
 
 # Thinc: A refreshing functional take on deep learning, compatible with your favorite libraries
 
-### From the makers of [spaCy](https://spacy.io), [Prodigy](https://prodi.gy) and [FastAPI](https://fastapi.tiangolo.com)
+### From the makers of [spaCy](https://spacy.io) and [Prodigy](https://prodi.gy)
 
 [Thinc](https://thinc.ai) is a **lightweight deep learning library** that offers an elegant,
 type-checked, functional-programming API for **composing models**, with support

azure-pipelines.yml

@@ -23,7 +23,7 @@ jobs:
         imageName: 'windows-2019'
         python.version: '3.6'
       Python37Mac:
-        imageName: 'macos-10.15'
+        imageName: 'macos-latest'
         python.version: '3.7'
       Python38Linux:
         imageName: 'ubuntu-latest'

requirements.txt

@@ -34,3 +34,4 @@ nbformat>=5.0.4,<5.2.0
 # Test to_disk/from_disk against pathlib.Path subclasses
 pathy>=0.3.5
 black>=22.0,<23.0
+confection>=0.0.1,<1.0.0

setup.cfg

@@ -45,6 +45,7 @@ install_requires =
     wasabi>=0.8.1,<1.1.0
     srsly>=2.4.0,<3.0.0
     catalogue>=2.0.4,<2.1.0
+    confection>=0.0.1,<1.0.0
     # Third-party dependencies
     setuptools
     numpy>=1.15.0

thinc/about.py

@@ -1,2 +1,2 @@
-__version__ = "8.1.0"
+__version__ = "8.1.1"
 __release__ = True

thinc/config.py

@@ -59,7 +59,7 @@
 from .list2padded import list2padded
 from .ragged2list import ragged2list
 from .padded2list import padded2list
-from .remap_ids import remap_ids
+from .remap_ids import remap_ids, remap_ids_v2
 from .strings2arrays import strings2arrays
 from .with_array import with_array
 from .with_array2d import with_array2d

thinc/layers/__init__.py

@@ -1,18 +1,23 @@
-from typing import Tuple, Callable, Sequence, Dict, Any
+from typing import Tuple, Callable, Sequence, cast
+from typing import Dict, Union, Optional, Hashable, Any
 
 from ..model import Model
 from ..config import registry
-from ..types import Ints2d, DTypes
+from ..types import Ints1d, Ints2d, DTypes
+from ..util import is_xp_array, to_numpy
 
 
-InT = Sequence[Any]
+InT = Union[Sequence[Hashable], Ints1d, Ints2d]
 OutT = Ints2d
 
+InT_v1 = Sequence[Any]
+OutT_v1 = Ints2d
+
 
 @registry.layers("remap_ids.v1")
 def remap_ids(
     mapping_table: Dict[Any, int] = {}, default: int = 0, dtype: DTypes = "i"
-) -> Model[InT, OutT]:
+) -> Model[InT_v1, OutT_v1]:
     """Remap string or integer inputs using a mapping table, usually as a
     preprocess before embeddings. The mapping table can be passed in on input,
     or updated after the layer has been created. The mapping table is stored in
@@ -26,7 +31,7 @@ def remap_ids(
 
 
 def forward(
-    model: Model[InT, OutT], inputs: InT, is_train: bool
+    model: Model[InT_v1, OutT_v1], inputs: InT_v1, is_train: bool
 ) -> Tuple[OutT, Callable]:
     table = model.attrs["mapping_table"]
     default = model.attrs["default"]
@@ -35,7 +40,60 @@ def forward(
     arr = model.ops.asarray2i(values, dtype=dtype)
     output = model.ops.reshape2i(arr, -1, 1)
 
-    def backprop(dY: OutT) -> InT:
+    def backprop(dY: OutT_v1) -> InT:
         return []
 
     return output, backprop
+
+
+@registry.layers("remap_ids.v2")
+def remap_ids_v2(
+    mapping_table: Optional[Union[Dict[int, int], Dict[str, int]]] = None,
+    default: int = 0,
+    *,
+    column: Optional[int] = None
+) -> Model[InT, OutT]:
+    """Remap string or integer inputs using a mapping table,
+    usually as a preprocessing step before embeddings.
+    The mapping table can be passed in on input,
+    or updated after the layer has been created.
+    The mapping table is stored in the "mapping_table" attribute.
+    Two dimensional arrays can be provided as input in which case
+    the 'column' chooses which column to process. This is useful
+    to work together with FeatureExtractor in spaCy.
+    """
+    return Model(
+        "remap_ids",
+        forward_v2,
+        attrs={"mapping_table": mapping_table, "default": default, "column": column},
+    )
+
+
+def forward_v2(
+    model: Model[InT, OutT], inputs: InT, is_train: bool
+) -> Tuple[OutT, Callable]:
+    table = model.attrs["mapping_table"]
+    if table is None:
+        raise ValueError("'mapping table' not set")
+    default = model.attrs["default"]
+    column = model.attrs["column"]
+    if is_xp_array(inputs):
+        xp_input = True
+        if column is not None:
+            idx = to_numpy(cast(Ints2d, inputs)[:, column])
+        else:
+            idx = to_numpy(inputs)
+    else:
+        xp_input = False
+        idx = inputs
+    values = [table.get(x, default) for x in idx]
+    arr = model.ops.asarray2i(values, dtype="i")
+    output = model.ops.reshape2i(arr, -1, 1)
+
+    def backprop(dY: OutT) -> InT:
+        if xp_input:
+            return model.ops.xp.empty(dY.shape)  # type: ignore
+        else:
+            return []
+
+    return output, backprop

thinc/layers/remap_ids.py

@@ -51,12 +51,11 @@ def __init__(
         self._backoff_factor = backoff_factor
         self._growth_interval = growth_interval
 
-        self._found_inf = torch.full((1,), 0.0)
         self._growth_tracker = torch.full((1,), 0, dtype=torch.int)
         self._scale = torch.full((1,), init_scale)
+        self._found_inf = False
 
     def to_(self, device):
-        self._found_inf = self._found_inf.to(device)
         self._growth_tracker = self._growth_tracker.to(device)
         self._scale = self._scale.to(device)
 
@@ -132,7 +131,7 @@ def _tensors_per_device(self, tensors):
 
     @property
     def found_inf(self):
-        return bool(self._found_inf) != 0
+        return self._found_inf
 
     def unscale(self, tensors):
         """Unscale the given tensors. Returns True if any of the gradients were infinite."""
@@ -152,9 +151,10 @@ def unscale(self, tensors):
                 device_tensors, found_inf_device, inv_scale_device
             )
 
-            self._found_inf += found_inf_device.to(self._found_inf.device)
+            if bool(found_inf_device != 0):
+                self._found_inf = True
 
-        return bool(self._found_inf != 0)
+        return self._found_inf
 
     def update(self):
         """
@@ -165,14 +165,17 @@ def update(self):
         if not self._enabled:
             return
 
+        found_inf_device = torch.full(
+            (1,), 1.0 if self._found_inf else 0.0, device=self._scale.device
+        )
         torch._amp_update_scale_(
             self._scale,
             self._growth_tracker,
-            self._found_inf,
+            found_inf_device,
             self._growth_factor,
             self._backoff_factor,
             self._growth_interval,
         )
 
         # Clear infinity found status
-        self._found_inf = torch.zeros_like(self._found_inf)
+        self._found_inf = False

thinc/shims/pytorch_grad_scaler.py

@@ -128,7 +128,8 @@ def assert_data_match(Y, out_data):
     # ("CauchySimilarity.v1", {}, (array2d, array2d), array1d),
     ("ParametricAttention.v1", {}, ragged, ragged),
     ("SparseLinear.v1", {}, (numpy.asarray([1, 2, 3], dtype="uint64"), array1d, numpy.asarray([1, 1], dtype="i")), array2d),
-    ("remap_ids.v1", {"dtype": "f"}, ["a", 1, 5.0], array2dint)
+    ("remap_ids.v1", {"dtype": "f"}, ["a", 1, 5.0], array2dint),
+    ("remap_ids.v2", {"mapping_table": {}, "column": 1}, numpy.array([[1, 2, 3], [4, 5, 6]]).T, array2dint)
     # fmt: on
 ]