urrgh

mittagessen · Apr 7, 2024 · 9ca85d6 · 9ca85d6
1 parent 46e71bb
commit 9ca85d6
Showing 1 changed file with 5 additions and 3 deletions.
diff --git 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: