infoset_saver.py

1importjson2importpathlib3fromtypingimportDict45fromlabmlimportexperiment6fromlabml\_nn.cfrimportInfoSet

9classInfoSetSaver(experiment.ModelSaver):

10def\_\_init\_\_(self,infosets:Dict[str,InfoSet]):11self.infosets=infosets

13defsave(self,checkpoint\_path:pathlib.Path)-\>any:14data={key:infoset.to\_dict()forkey,infosetinself.infosets.items()}15file\_name=f"infosets.json"1617withopen(str(checkpoint\_path/file\_name),'w')asf:18f.write(json.dumps(data))1920returnfile\_name

22defload(self,checkpoint\_path:pathlib.Path,file\_name:str):23withopen(str(checkpoint\_path/file\_name),'w')asf:24data=json.loads(f.read())2526forkey,dindata.items():27self.infosets[key]=InfoSet.from\_dict(d)