More general remap_ids

thinc/layers/remap_ids.py

@@ -1,32 +1,37 @@
-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(
+def remap_ids_v1(
     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
     the "mapping_table" attribute.
     """
     return Model(
         "remap_ids",
-        forward,
+        forward_v1,
         attrs={"mapping_table": mapping_table, "dtype": dtype, "default": default},
     )
 
 
-def forward(
-    model: Model[InT, OutT], inputs: InT, is_train: bool
+def forward_v1(
+    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(
+    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,
+        attrs={"mapping_table": mapping_table, "default": default, "column": column},
+    )
+
+
+def forward(
+    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/tests/layers/test_layers_api.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
 ]
 

website/docs/api-layers.md

@@ -1267,21 +1267,25 @@ https://github.com/explosion/thinc/blob/master/thinc/layers/padded2list.py
 
 <inline-list>
 
-- **Input:** <tt>Sequence[Any]</tt>
+- **Input:** <tt>Union[Sequence[Hashable], Ints1d, Ints2d]</tt>
 - **Output:** <ndarray>Ints2d</ndarray>
 
 </inline-list>
 
-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 the `"mapping_table"`
-attribute.
+Remap a sequence of strings, integers or other Hashable inputs using a
+mapping table, usually as a preprocessing step before embeddings. The
+input can also be a two dimensional integer array in which case the
+`column: int` attribute tells the `remap_ids` layer which column of the
+array to map. For 2D inputs the `column` has to be set on initialization.
+The mapping table can be set at initialization, but can also be passed in on input, 
+When provided on initializtion the mapping table is stored in the `"mapping_table"`
+attribute. 
 
 | Argument        | Type                                  | Description                                                                                                  |
 | --------------- | ------------------------------------- | ------------------------------------------------------------------------------------------------------------ |
 | `mapping_table` | <tt>Dict[Any, int]</tt>               | The mapping table to use. Can also be set after initialization by writing to `model.attrs["mapping_table"]`. |
 | `default`       | <tt>int</tt>                          | The default value if the input does not have an entry in the mapping table.                                  |
-| `dtype`         | <tt>DTypes</tt>                       | The data type of the array.                                                                                  |
+| `column`         | <tt>int</tt>                       | The column to apply the mapper to in case of 2D input.                                                                                  |
 | **RETURNS**     | <tt>Model[Sequence[Any], Ints2d]</tt> | The layer to compute the transformation.                                                                     |
 
 ```python