Added support for RobertaModel to transformers.py (#1864)

* Add support for RobertaModel to transformers.py

* Adding blank-lines to fix build failure, E302 expected 2 blank lines, found 1

* Fix missing parentheses an ...if...else... statements

---------

Co-authored-by: Minh-Thuc <[email protected]>

Author: sheldonrobinson

python/ctranslate2/converters/opennmt_tf.py

@@ -291,9 +291,11 @@ def set_multi_head_attention(self, spec, module, self_attention=False):
     def set_layer_norm_from_wrapper(self, spec, module):
         self.set_layer_norm(
             spec,
-            module.output_layer_norm
-            if module.input_layer_norm is None
-            else module.input_layer_norm,
+            (
+                module.output_layer_norm
+                if module.input_layer_norm is None
+                else module.input_layer_norm
+            ),
         )
 
     def set_layer_norm(self, spec, module):

python/ctranslate2/converters/transformers.py

@@ -2667,6 +2667,170 @@ def set_position_encodings(self, spec, module):
             spec.encodings = spec.encodings[offset + 1 :]
 
 
+@register_loader("RobertaConfig")
+class RobertaLoader(ModelLoader):
+    @property
+    def architecture_name(self):
+        return "RobertaModel"
+
+    def get_model_spec(self, model):
+        assert model.config.position_embedding_type == "absolute"
+
+        encoder_spec = transformer_spec.TransformerEncoderSpec(
+            model.config.num_hidden_layers,
+            model.config.num_attention_heads,
+            pre_norm=False,
+            activation=_SUPPORTED_ACTIVATIONS[model.config.hidden_act],
+            layernorm_embedding=True,
+            num_source_embeddings=2,
+            embeddings_merge=common_spec.EmbeddingsMerge.ADD,
+        )
+
+        if model.pooler is None:
+            pooling_layer = False
+        else:
+            pooling_layer = True
+
+        spec = transformer_spec.TransformerEncoderModelSpec(
+            encoder_spec,
+            pooling_layer=pooling_layer,
+            pooling_activation=common_spec.Activation.Tanh,
+        )
+
+        spec.encoder.scale_embeddings = False
+
+        self.set_embeddings(
+            spec.encoder.embeddings[0], model.embeddings.word_embeddings
+        )
+        self.set_embeddings(
+            spec.encoder.embeddings[1], model.embeddings.token_type_embeddings
+        )
+        self.set_position_encodings(
+            spec.encoder.position_encodings,
+            model.embeddings.position_embeddings,
+        )
+        self.set_layer_norm(
+            spec.encoder.layernorm_embedding, model.embeddings.LayerNorm
+        )
+        if pooling_layer:
+            self.set_linear(spec.pooler_dense, model.pooler.dense)
+
+        for layer_spec, layer in zip(spec.encoder.layer, model.encoder.layer):
+            split_layers = [common_spec.LinearSpec() for _ in range(3)]
+            self.set_linear(split_layers[0], layer.attention.self.query)
+            self.set_linear(split_layers[1], layer.attention.self.key)
+            self.set_linear(split_layers[2], layer.attention.self.value)
+            utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers)
+
+            self.set_linear(
+                layer_spec.self_attention.linear[1], layer.attention.output.dense
+            )
+            self.set_layer_norm(
+                layer_spec.self_attention.layer_norm, layer.attention.output.LayerNorm
+            )
+
+            self.set_linear(layer_spec.ffn.linear_0, layer.intermediate.dense)
+            self.set_linear(layer_spec.ffn.linear_1, layer.output.dense)
+            self.set_layer_norm(layer_spec.ffn.layer_norm, layer.output.LayerNorm)
+
+        return spec
+
+    def set_vocabulary(self, spec, tokens):
+        spec.register_vocabulary(tokens)
+
+    def set_config(self, config, model, tokenizer):
+        config.unk_token = tokenizer.unk_token
+        config.layer_norm_epsilon = model.config.layer_norm_eps
+
+    def set_position_encodings(self, spec, module):
+        spec.encodings = module.weight
+        offset = getattr(module, "padding_idx", 0)
+        if offset > 0:
+            spec.encodings = spec.encodings[offset + 1 :]
+
+
+@register_loader("CamembertConfig")
+class CamembertLoader(ModelLoader):
+    @property
+    def architecture_name(self):
+        return "CamembertModel"
+
+    def get_model_spec(self, model):
+        assert model.config.position_embedding_type == "absolute"
+
+        encoder_spec = transformer_spec.TransformerEncoderSpec(
+            model.config.num_hidden_layers,
+            model.config.num_attention_heads,
+            pre_norm=False,
+            activation=_SUPPORTED_ACTIVATIONS[model.config.hidden_act],
+            layernorm_embedding=True,
+            num_source_embeddings=2,
+            embeddings_merge=common_spec.EmbeddingsMerge.ADD,
+        )
+
+        if model.pooler is None:
+            pooling_layer = False
+        else:
+            pooling_layer = True
+
+        spec = transformer_spec.TransformerEncoderModelSpec(
+            encoder_spec,
+            pooling_layer=pooling_layer,
+            pooling_activation=common_spec.Activation.Tanh,
+        )
+
+        spec.encoder.scale_embeddings = False
+
+        self.set_embeddings(
+            spec.encoder.embeddings[0], model.embeddings.word_embeddings
+        )
+        self.set_embeddings(
+            spec.encoder.embeddings[1], model.embeddings.token_type_embeddings
+        )
+        self.set_position_encodings(
+            spec.encoder.position_encodings,
+            model.embeddings.position_embeddings,
+        )
+        self.set_layer_norm(
+            spec.encoder.layernorm_embedding, model.embeddings.LayerNorm
+        )
+        if pooling_layer:
+            self.set_linear(spec.pooler_dense, model.pooler.dense)
+
+        for layer_spec, layer in zip(spec.encoder.layer, model.encoder.layer):
+            split_layers = [common_spec.LinearSpec() for _ in range(3)]
+            self.set_linear(split_layers[0], layer.attention.self.query)
+            self.set_linear(split_layers[1], layer.attention.self.key)
+            self.set_linear(split_layers[2], layer.attention.self.value)
+            utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers)
+
+            self.set_linear(
+                layer_spec.self_attention.linear[1], layer.attention.output.dense
+            )
+            self.set_layer_norm(
+                layer_spec.self_attention.layer_norm, layer.attention.output.LayerNorm
+            )
+
+            self.set_linear(layer_spec.ffn.linear_0, layer.intermediate.dense)
+            self.set_linear(layer_spec.ffn.linear_1, layer.output.dense)
+            self.set_layer_norm(layer_spec.ffn.layer_norm, layer.output.LayerNorm)
+
+        return spec
+
+    def set_vocabulary(self, spec, tokens):
+        spec.register_vocabulary(tokens)
+
+    def set_config(self, config, model, tokenizer):
+        config.unk_token = tokenizer.unk_token
+        config.layer_norm_epsilon = model.config.layer_norm_eps
+
+    def set_position_encodings(self, spec, module):
+        spec.encodings = module.weight
+        offset = getattr(module, "padding_idx", 0)
+        if offset > 0:
+            spec.encodings = spec.encodings[offset + 1 :]
+
+
 def main():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter

python/ctranslate2/converters/utils.py

@@ -25,9 +25,11 @@ def fuse_linear(spec, layers):
     if bias_dtype is not None:
         spec.bias = concatenate(
             [
-                layer.bias
-                if layer.has_bias()
-                else zeros([layer.weight.shape[0]], dtype=bias_dtype)
+                (
+                    layer.bias
+                    if layer.has_bias()
+                    else zeros([layer.weight.shape[0]], dtype=bias_dtype)
+                )
                 for layer in layers
             ]
         )