Source code for src.ml_helper.model

import uuid
from datetime import datetime

import sqlalchemy as db
from sqlalchemy_utils import UUIDType, create_database, database_exists
from sqlalchemy.orm import relationship, sessionmaker
from sqlalchemy.ext.declarative import declarative_base


[docs]ENGINE_URL = "postgres://user:password@localhost:5432/ml_helper"

[docs]metadata = db.MetaData()
[docs]Base = declarative_base(metadata=metadata)


[docs]class Model(Base):
[docs]    __tablename__ = "model"

[docs]    uuid = db.Column(UUIDType, primary_key=True, default=uuid.uuid4)

[docs]    id = db.Column(db.String, unique=True)
[docs]    last_updated = db.Column(db.DateTime,
                             default=datetime.now,
                             onupdate=datetime.now)

[docs]    params = db.Column(db.JSON)

[docs]    epoch = relationship("Epoch")


[docs]class Epoch(Base):
[docs]    __tablename__ = "epoch"
[docs]    __table_args__ = (
        db.UniqueConstraint("model_id", "number", name="unique_epoch"),
    )

[docs]    uuid = db.Column(UUIDType, primary_key=True, default=uuid.uuid4)
[docs]    model_id = db.Column(
        db.String,
        db.ForeignKey('model.id'),
    )

[docs]    created_at = db.Column(db.DateTime, default=datetime.now)

[docs]    number = db.Column(db.Integer, nullable=False)

[docs]    training_loss = db.Column(db.Float)
[docs]    eval_loss = db.Column(db.Float)

[docs]    training_F1 = db.Column(db.Float)
[docs]    eval_F1 = db.Column(db.Float)

[docs]    training_acc = db.Column(db.Float)
[docs]    eval_acc = db.Column(db.Float)


[docs]engine = db.create_engine(ENGINE_URL)
[docs]Session = sessionmaker(bind=engine)


[docs]def init_db():
    if not database_exists(ENGINE_URL):
        create_database(ENGINE_URL)
    metadata.drop_all(engine)
    metadata.create_all(engine)
    print("DB initialized.")


if __name__ == "__main__":
    init_db()