From 9ca85d6dc37cada1e9bdf8b8bbf65a605028fe21 Mon Sep 17 00:00:00 2001
From: Benjamin Kiessling <mittagessen@l.unchti.me>
Date: Sun, 7 Apr 2024 21:53:02 +0200
Subject: [PATCH] urrgh

---
 kraken/lib/vgsl.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

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: