Model.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. import json
  2. from pathlib import Path
  3. from pycs import db
  4. from pycs.database.base import NamedBaseModel
  5. class Model(NamedBaseModel):
  6. """
  7. database class for ML Models
  8. """
  9. description = db.Column(db.String)
  10. root_folder = db.Column(db.String, nullable=False, unique=True)
  11. supports_encoded = db.Column(db.String, nullable=False)
  12. # relationships to other models
  13. projects = db.relationship("Project", backref="model", lazy="dynamic")
  14. serialize_rules = ('-projects',)
  15. def serialize(self):
  16. result = super().serialize()
  17. result["supports"] = self.supports
  18. return result
  19. @property
  20. def supports(self):
  21. return json.loads(self.supports_encoded)
  22. @supports.setter
  23. def supports(self, value):
  24. if isinstance(value, str):
  25. self.supports_encoded = value
  26. elif isinstance(value, (dict, list)):
  27. self.supports_encoded = json.dumps(value)
  28. else:
  29. raise ValueError(f"Not supported type: {type(value)}")
  30. def copy_to(self, new_name: str, new_root_folder: str):
  31. model, is_new = Model.get_or_create(root_folder=new_root_folder)
  32. model.name = new_name
  33. model.description = self.description
  34. model.supports_encoded = self.supports_encoded
  35. self.commit()
  36. return model, is_new
  37. @classmethod
  38. def discover(cls, root: Path, config_name: str = "configuration.json"):
  39. for folder in Path(root).glob("*"):
  40. with open(folder / config_name) as f:
  41. config = json.load(f)
  42. # extract data
  43. name = config['name']
  44. description = config.get('description', None)
  45. supports = config['supports']
  46. model, _ = cls.get_or_create(root_folder=str(folder))
  47. model.name = name
  48. model.description = description
  49. model.supports = supports
  50. db.session.commit()