Bring back curated_encoder prefix

spacy_curated_transformers/models/architectures.py

@@ -3,6 +3,7 @@
 from typing import Any, Callable, List, Optional, Tuple, Union, cast
 
 import torch
+import torch.nn as nn
 from curated_transformers.layers import Activation, AttentionMask
 from curated_transformers.models import (
  ALBERTConfig,
@@ -1181,7 +1182,7 @@ def _pytorch_encoder(
  grad_scaler_config["enabled"] = mixed_precision
 
  model = PyTorchWrapper_v2(
- encoder,
+ _EncoderWrapper(encoder),
  convert_inputs=partial(
  _convert_inputs,
  max_model_seq_len=model_max_length,
@@ -1311,7 +1312,7 @@ def build_pytorch_checkpoint_loader_v2(*, path: Path) -> Callable[
 
  def load(model, X=None, Y=None):
  device = get_torch_default_device()
- encoder = model.shims[0]._model
+ encoder = model.shims[0]._model.curated_encoder
  assert isinstance(encoder, FromHFHub)
  fs = LocalFileSystem()
  encoder.from_fsspec_(fs=fs, model_path=path, device=device)
@@ -1325,3 +1326,16 @@ def _torch_dtype_from_str(dtype_as_str: str):
  if not isinstance(dtype, torch.dtype):
  raise ValueError(f"Invalid torch dtype `{dtype_as_str}`")
  return dtype
+
+
+class _EncoderWrapper(nn.Module):
+ """Small wrapper to add a prefix that can be used by eg. learning rate
+ schedules.
+ """
+
+ def __init__(self, encoder: nn.Module):
+ super().__init__()
+ self.curated_encoder = encoder
+
+ def forward(self, *args, **kwargs):
+ return self.curated_encoder.forward(*args, **kwargs)

spacy_curated_transformers/models/hf_loader.py

@@ -24,7 +24,7 @@ def build_hf_transformer_encoder_loader_v1(
  """
 
  def load(model, X=None, Y=None):
- encoder = model.shims[0]._model
+ encoder = model.shims[0]._model.curated_encoder
  assert isinstance(encoder, FromHFHub)
  device = model.shims[0].device
  encoder.from_hf_hub_(name=name, revision=revision, device=device)

spacy_curated_transformers/tests/conftest.py

@@ -40,7 +40,6 @@ def getopt(opt):
 
 @pytest.fixture
 def test_dir(request):
- print(request.fspath)
  return Path(request.fspath).parent
 
 

spacy_curated_transformers/tests/models/test_transformer_model.py

@@ -178,3 +178,44 @@ def test_pytorch_checkpoint_loader(test_config):
  "spacy-curated-transformers.PyTorchCheckpointLoader.v1"
  )(path=Path(checkpoint_path))
  model.initialize()
+
+
+@pytest.mark.parametrize(
+ "test_config",
+ [
+ (
+ build_albert_transformer_model_v1,
+ build_sentencepiece_encoder_v1(),
+ 1000,
+ ),
+ (
+ build_bert_transformer_model_v1,
+ build_bert_wordpiece_encoder_v1(),
+ 1000,
+ ),
+ (
+ build_roberta_transformer_model_v1,
+ build_byte_bpe_encoder_v1(),
+ 1000,
+ ),
+ ],
+)
+def test_encoder_prefix(test_config):
+ model_factory, piece_encoder, vocab_size = test_config
+
+ # Curated Transformers needs the config to get the model hyperparameters.
+ with_spans = build_with_strided_spans_v1(stride=96, window=128)
+ model = model_factory(
+ hidden_width=32,
+ intermediate_width=37,
+ num_hidden_layers=4,
+ num_attention_heads=4,
+ piece_encoder=piece_encoder,
+ vocab_size=vocab_size,
+ with_spans=with_spans,
+ )
+
+ for name, _ in model.get_ref("transformer").shims[0]._model.named_parameters():
+ assert name.startswith(
+ "curated_encoder."
+ ), f"Parameter name '{name} does not start with 'curated_encoder'"