diff --git a/alembic/versions/0003_fc2b_ml_pipeline.py b/alembic/versions/0003_fc2b_ml_pipeline.py new file mode 100644 index 0000000..584bffe --- /dev/null +++ b/alembic/versions/0003_fc2b_ml_pipeline.py @@ -0,0 +1,172 @@ +"""fc2b: ML pipeline — allowlist, aliases, centroids, ml_settings + +Revision ID: 0003 +Revises: 0002 +Create Date: 2026-05-15 + +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op +from pgvector.sqlalchemy import Vector + +revision: str = "0003" +down_revision: Union[str, None] = "0002" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # 3.1 rename wd14_* -> tagger_* + op.alter_column("image_record", "wd14_predictions", new_column_name="tagger_predictions") + op.alter_column( + "image_record", "wd14_model_version", new_column_name="tagger_model_version" + ) + + # 3.2 tag_allowlist + op.create_table( + "tag_allowlist", + sa.Column("tag_id", sa.Integer(), nullable=False), + sa.Column( + "min_confidence", sa.Float(), nullable=False, server_default="0.95" + ), + sa.Column( + "added_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.ForeignKeyConstraint( + ["tag_id"], ["tag.id"], name="fk_tag_allowlist_tag_id_tag", + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("tag_id", name="pk_tag_allowlist"), + sa.CheckConstraint( + "min_confidence > 0 AND min_confidence <= 1", + name="ck_tag_allowlist_confidence_range", + ), + ) + + # 3.3 tag_suggestion_rejection + op.create_table( + "tag_suggestion_rejection", + sa.Column("image_record_id", sa.Integer(), nullable=False), + sa.Column("tag_id", sa.Integer(), nullable=False), + sa.Column( + "rejected_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.ForeignKeyConstraint( + ["image_record_id"], ["image_record.id"], + name="fk_tsr_image_record_id_image_record", ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["tag_id"], ["tag.id"], name="fk_tsr_tag_id_tag", ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint( + "image_record_id", "tag_id", name="pk_tag_suggestion_rejection" + ), + ) + op.create_index( + "ix_tag_suggestion_rejection_tag", "tag_suggestion_rejection", ["tag_id"] + ) + + # 3.4 tag_alias + op.create_table( + "tag_alias", + sa.Column("alias_string", sa.String(length=255), nullable=False), + sa.Column("alias_category", sa.String(length=32), nullable=False), + sa.Column("canonical_tag_id", sa.Integer(), nullable=False), + sa.Column( + "created_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.ForeignKeyConstraint( + ["canonical_tag_id"], ["tag.id"], + name="fk_tag_alias_canonical_tag_id_tag", ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint( + "alias_string", "alias_category", name="pk_tag_alias" + ), + ) + op.create_index("ix_tag_alias_canonical", "tag_alias", ["canonical_tag_id"]) + + # 3.5 tag_reference_embedding (centroids) + op.create_table( + "tag_reference_embedding", + sa.Column("tag_id", sa.Integer(), nullable=False), + sa.Column("embedding", Vector(1152), nullable=False), + sa.Column("reference_count", sa.Integer(), nullable=False), + sa.Column("model_version", sa.String(length=128), nullable=False), + sa.Column( + "updated_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.ForeignKeyConstraint( + ["tag_id"], ["tag.id"], + name="fk_tag_reference_embedding_tag_id_tag", ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("tag_id", name="pk_tag_reference_embedding"), + ) + + # 3.6 ml_settings singleton + op.create_table( + "ml_settings", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "suggestion_threshold_artist", sa.Float(), nullable=False, + server_default="0.30", + ), + sa.Column( + "suggestion_threshold_character", sa.Float(), nullable=False, + server_default="0.50", + ), + sa.Column( + "suggestion_threshold_copyright", sa.Float(), nullable=False, + server_default="0.50", + ), + sa.Column( + "suggestion_threshold_general", sa.Float(), nullable=False, + server_default="0.95", + ), + sa.Column( + "centroid_similarity_threshold", sa.Float(), nullable=False, + server_default="0.55", + ), + sa.Column( + "min_reference_images", sa.Integer(), nullable=False, + server_default="5", + ), + sa.Column( + "tagger_model_version", sa.String(length=128), nullable=False, + server_default="camie-tagger-v2", + ), + sa.Column( + "embedder_model_version", sa.String(length=128), nullable=False, + server_default="siglip-so400m-patch14-384", + ), + sa.Column( + "updated_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.PrimaryKeyConstraint("id", name="pk_ml_settings"), + sa.CheckConstraint("id = 1", name="ck_ml_settings_singleton"), + ) + op.execute("INSERT INTO ml_settings (id) VALUES (1)") + + +def downgrade() -> None: + op.drop_table("ml_settings") + op.drop_table("tag_reference_embedding") + op.drop_index("ix_tag_alias_canonical", table_name="tag_alias") + op.drop_table("tag_alias") + op.drop_index( + "ix_tag_suggestion_rejection_tag", table_name="tag_suggestion_rejection" + ) + op.drop_table("tag_suggestion_rejection") + op.drop_table("tag_allowlist") + op.alter_column( + "image_record", "tagger_model_version", new_column_name="wd14_model_version" + ) + op.alter_column( + "image_record", "tagger_predictions", new_column_name="wd14_predictions" + ) diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 8052e9b..69e4e73 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -9,9 +9,14 @@ from .image_record import ImageRecord from .import_batch import ImportBatch from .import_settings import ImportSettings from .import_task import ImportTask +from .ml_settings import MLSettings from .post import Post from .source import Source from .tag import Tag, TagKind, image_tag +from .tag_alias import TagAlias +from .tag_allowlist import TagAllowlist +from .tag_reference_embedding import TagReferenceEmbedding +from .tag_suggestion_rejection import TagSuggestionRejection __all__ = [ "Base", @@ -28,4 +33,9 @@ __all__ = [ "ImportBatch", "ImportTask", "ImportSettings", + "MLSettings", + "TagAlias", + "TagAllowlist", + "TagReferenceEmbedding", + "TagSuggestionRejection", ] diff --git a/backend/app/models/image_record.py b/backend/app/models/image_record.py index 9df4845..b303e7c 100644 --- a/backend/app/models/image_record.py +++ b/backend/app/models/image_record.py @@ -56,8 +56,8 @@ class ImageRecord(Base): ) # ML fields (populated by FC-2's ml-worker) - wd14_predictions: Mapped[dict | None] = mapped_column(JSON, nullable=True) - wd14_model_version: Mapped[str | None] = mapped_column(String(128), nullable=True) + tagger_predictions: Mapped[dict | None] = mapped_column(JSON, nullable=True) + tagger_model_version: Mapped[str | None] = mapped_column(String(128), nullable=True) # 1152 = SigLIP-so400m embedding dim. Swapping models in FC-2 may require # a column-width migration. siglip_embedding: Mapped[list[float] | None] = mapped_column(Vector(1152), nullable=True) diff --git a/backend/app/models/ml_settings.py b/backend/app/models/ml_settings.py new file mode 100644 index 0000000..9586b8c --- /dev/null +++ b/backend/app/models/ml_settings.py @@ -0,0 +1,42 @@ +"""MLSettings — single-row table holding ML pipeline tunables.""" + +from datetime import datetime + +from sqlalchemy import CheckConstraint, DateTime, Float, Integer, String, func +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class MLSettings(Base): + __tablename__ = "ml_settings" + __table_args__ = (CheckConstraint("id = 1", name="ck_ml_settings_singleton"),) + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + suggestion_threshold_artist: Mapped[float] = mapped_column( + Float, nullable=False, default=0.30 + ) + suggestion_threshold_character: Mapped[float] = mapped_column( + Float, nullable=False, default=0.50 + ) + suggestion_threshold_copyright: Mapped[float] = mapped_column( + Float, nullable=False, default=0.50 + ) + suggestion_threshold_general: Mapped[float] = mapped_column( + Float, nullable=False, default=0.95 + ) + centroid_similarity_threshold: Mapped[float] = mapped_column( + Float, nullable=False, default=0.55 + ) + min_reference_images: Mapped[int] = mapped_column( + Integer, nullable=False, default=5 + ) + tagger_model_version: Mapped[str] = mapped_column( + String(128), nullable=False, default="camie-tagger-v2" + ) + embedder_model_version: Mapped[str] = mapped_column( + String(128), nullable=False, default="siglip-so400m-patch14-384" + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) diff --git a/backend/app/models/tag_alias.py b/backend/app/models/tag_alias.py new file mode 100644 index 0000000..3a21375 --- /dev/null +++ b/backend/app/models/tag_alias.py @@ -0,0 +1,24 @@ +"""TagAlias — maps a model's (name, category) prediction to the operator's +canonical tag. Resolved at suggestion-read time so raw predictions stay +unmolested in image_record.tagger_predictions. +""" + +from datetime import datetime + +from sqlalchemy import DateTime, ForeignKey, String, func +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class TagAlias(Base): + __tablename__ = "tag_alias" + + alias_string: Mapped[str] = mapped_column(String(255), primary_key=True) + alias_category: Mapped[str] = mapped_column(String(32), primary_key=True) + canonical_tag_id: Mapped[int] = mapped_column( + ForeignKey("tag.id", ondelete="CASCADE"), nullable=False, index=True + ) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) diff --git a/backend/app/models/tag_allowlist.py b/backend/app/models/tag_allowlist.py new file mode 100644 index 0000000..6508178 --- /dev/null +++ b/backend/app/models/tag_allowlist.py @@ -0,0 +1,26 @@ +"""TagAllowlist — tags the operator opted-in to auto-apply via maintenance.""" + +from datetime import datetime + +from sqlalchemy import CheckConstraint, DateTime, Float, ForeignKey, func +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class TagAllowlist(Base): + __tablename__ = "tag_allowlist" + __table_args__ = ( + CheckConstraint( + "min_confidence > 0 AND min_confidence <= 1", + name="ck_tag_allowlist_confidence_range", + ), + ) + + tag_id: Mapped[int] = mapped_column( + ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True + ) + min_confidence: Mapped[float] = mapped_column(Float, nullable=False, default=0.95) + added_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) diff --git a/backend/app/models/tag_reference_embedding.py b/backend/app/models/tag_reference_embedding.py new file mode 100644 index 0000000..2dc841b --- /dev/null +++ b/backend/app/models/tag_reference_embedding.py @@ -0,0 +1,23 @@ +"""TagReferenceEmbedding — per-tag centroid (mean SigLIP embedding of members).""" + +from datetime import datetime + +from pgvector.sqlalchemy import Vector +from sqlalchemy import DateTime, ForeignKey, Integer, String, func +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class TagReferenceEmbedding(Base): + __tablename__ = "tag_reference_embedding" + + tag_id: Mapped[int] = mapped_column( + ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True + ) + embedding: Mapped[list[float]] = mapped_column(Vector(1152), nullable=False) + reference_count: Mapped[int] = mapped_column(Integer, nullable=False) + model_version: Mapped[str] = mapped_column(String(128), nullable=False) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) diff --git a/backend/app/models/tag_suggestion_rejection.py b/backend/app/models/tag_suggestion_rejection.py new file mode 100644 index 0000000..8a834fe --- /dev/null +++ b/backend/app/models/tag_suggestion_rejection.py @@ -0,0 +1,25 @@ +"""TagSuggestionRejection — per-image dismissed suggestions. + +Prevents re-suggestion AND prevents allowlist auto-apply on that image. +""" + +from datetime import datetime + +from sqlalchemy import DateTime, ForeignKey, func +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class TagSuggestionRejection(Base): + __tablename__ = "tag_suggestion_rejection" + + image_record_id: Mapped[int] = mapped_column( + ForeignKey("image_record.id", ondelete="CASCADE"), primary_key=True + ) + tag_id: Mapped[int] = mapped_column( + ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True, index=True + ) + rejected_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) diff --git a/tests/test_migration_0003.py b/tests/test_migration_0003.py new file mode 100644 index 0000000..a46f76c --- /dev/null +++ b/tests/test_migration_0003.py @@ -0,0 +1,56 @@ +"""Smoke test for migration 0003: model classes import, schema shape correct.""" + +from backend.app.models import ( + Base, + ImageRecord, + MLSettings, + TagAlias, + TagAllowlist, + TagReferenceEmbedding, + TagSuggestionRejection, +) + + +def test_new_tables_registered(): + expected = { + "tag_allowlist", + "tag_suggestion_rejection", + "tag_alias", + "tag_reference_embedding", + "ml_settings", + } + assert expected.issubset(Base.metadata.tables.keys()) + + +def test_image_record_columns_renamed(): + cols = {c.name for c in ImageRecord.__table__.columns} + assert "tagger_predictions" in cols + assert "tagger_model_version" in cols + assert "wd14_predictions" not in cols + assert "wd14_model_version" not in cols + + +def test_tag_alias_composite_pk(): + pk_cols = {c.name for c in TagAlias.__table__.primary_key.columns} + assert pk_cols == {"alias_string", "alias_category"} + + +def test_ml_settings_singleton_constraint(): + names = {c.name for c in MLSettings.__table__.constraints} + assert "ck_ml_settings_singleton" in names + + +def test_tag_reference_embedding_has_vector(): + cols = {c.name for c in TagReferenceEmbedding.__table__.columns} + assert "embedding" in cols + assert "reference_count" in cols + + +def test_tag_allowlist_confidence_check(): + names = {c.name for c in TagAllowlist.__table__.constraints} + assert "ck_tag_allowlist_confidence_range" in names + + +def test_tag_suggestion_rejection_pk(): + pk_cols = {c.name for c in TagSuggestionRejection.__table__.primary_key.columns} + assert pk_cols == {"image_record_id", "tag_id"}