diff --git a/kraken/lib/vgsl.py b/kraken/lib/vgsl.py index 52444d4e5..19238214b 100644 --- a/kraken/lib/vgsl.py +++ b/kraken/lib/vgsl.py @@ -2,6 +2,8 @@ VGSL plumbing """ import json +import base64 + import logging import re from os import PathLike @@ -307,7 +309,7 @@ def _deserialize_layers(name, layer): if 'codec' in mlmodel.user_defined_metadata: codec = json.loads(mlmodel.user_defined_metadata['codec']) if codec['type'] == 'SentencePiece': - nn.add_codec(SentencePieceCodec(mlmodel.user_defined_metadata['sp_codec'])) + nn.add_codec(SentencePieceCodec(base64.b64decode(codec['spp']))) nn.user_metadata: Dict[str, Any] = {'accuracy': [], 'metrics': [], @@ -411,8 +413,8 @@ def _serialize_layer(net, input, net_builder): mlmodel.short_description = 'kraken model' mlmodel.user_defined_metadata['vgsl'] = '[' + ' '.join(self.named_spec) + ']' if self.codec: - mlmodel.user_defined_metadata['codec'] = json.dumps({'type': 'SentencePiece'}) - mlmodel.user_defined_metadata['sp_codec'] = self.codec.spp.serialized_model_proto() + mlmodel.user_defined_metadata['codec'] = json.dumps({'type': 'SentencePiece', + 'spp': base64.b64encode(self.codec.spp.serialized_model_proto())}) if self.user_metadata: mlmodel.user_defined_metadata['kraken_meta'] = json.dumps(self.user_metadata) if self.aux_layers: