Import message trees from jsonl file (LAION-AI#964)

* add new backlog_ranking tree state

* add first version of import script

* allow activation of trees during import

* add min_active_rankings_per_lang config param

* add settings docstring

Author: andreaskoepf

backend/alembic/versions/2023_01_28_1157-49d8445b4c90_add_origin_column_to_message_tree_state.py

@@ -0,0 +1,31 @@
+"""add origin column to message_tree_state
+
+Revision ID: 49d8445b4c90
+Revises: f856bf19d32b
+Create Date: 2023-01-28 11:57:45.580027
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "49d8445b4c90"
+down_revision = "f856bf19d32b"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("message", sa.Column("synthetic", sa.Boolean(), server_default=sa.text("false"), nullable=False))
+    op.add_column("message", sa.Column("model_name", sa.String(length=1024), nullable=True))
+    op.add_column("message_tree_state", sa.Column("origin", sa.String(length=1024), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("message_tree_state", "origin")
+    op.drop_column("message", "model_name")
+    op.drop_column("message", "synthetic")
+    # ### end Alembic commands ###

backend/import.py

@@ -0,0 +1,187 @@
+import argparse
+import json
+from pathlib import Path
+from typing import Optional
+from uuid import UUID
+
+import oasst_backend.models.db_payload as db_payload
+import oasst_backend.utils.database_utils as db_utils
+import pydantic
+from loguru import logger
+from oasst_backend.api.deps import create_api_client
+from oasst_backend.models import ApiClient, Message
+from oasst_backend.models.message_tree_state import MessageTreeState
+from oasst_backend.models.message_tree_state import State as TreeState
+from oasst_backend.models.payload_column_type import PayloadContainer
+from oasst_backend.prompt_repository import PromptRepository
+from oasst_backend.user_repository import UserRepository
+from oasst_backend.utils.tree_export import ExportMessageNode, ExportMessageTree
+from sqlmodel import Session
+
+# well known id
+IMPORT_API_CLIENT_ID = UUID("bd8fde8b-1d8e-4e9a-9966-e96d000f8363")
+
+
+class Importer:
+    def __init__(self, db: Session, origin: str, model_name: Optional[str] = None):
+        self.db = db
+        self.origin = origin
+        self.model_name = model_name
+
+        # get import api client
+        api_client = db.query(ApiClient).filter(ApiClient.id == IMPORT_API_CLIENT_ID).first()
+        if not api_client:
+            api_client = create_api_client(
+                session=db,
+                description="API client used for importing data",
+                frontend_type="import",
+                force_id=IMPORT_API_CLIENT_ID,
+            )
+
+        ur = UserRepository(db, api_client)
+        self.import_user = ur.lookup_system_user(username="import")
+        self.pr = PromptRepository(db=db, api_client=api_client, user_repository=ur)
+        self.api_client = api_client
+
+    def fetch_message_tree_state(self, message_tree_id: UUID) -> MessageTreeState:
+        return self.db.query(MessageTreeState).filter(MessageTreeState.message_tree_id == message_tree_id).one_or_none()
+
+    def import_message(
+        self, message: ExportMessageNode, message_tree_id: UUID, parent_id: Optional[UUID] = None
+    ) -> Message:
+        payload = db_payload.MessagePayload(text=message.text)
+        msg = Message(
+            id=message.message_id,
+            message_tree_id=message_tree_id,
+            frontend_message_id=message.message_id,
+            parent_id=parent_id,
+            review_count=message.review_count or 0,
+            lang=message.lang or "en",
+            review_result=True,
+            synthetic=message.synthetic if message.synthetic is not None else True,
+            model_name=message.model_name or self.model_name,
+            role=message.role,
+            api_client_id=self.api_client.id,
+            payload_type=type(payload).__name__,
+            payload=PayloadContainer(payload=payload),
+            user_id=self.import_user.id,
+        )
+        self.db.add(msg)
+        if message.replies:
+            for r in message.replies:
+                self.import_message(r, message_tree_id=message_tree_id, parent_id=msg.id)
+        self.db.flush()
+        if parent_id is None:
+            self.pr.update_children_counts(msg.id)
+        self.db.refresh(msg)
+        return msg
+
+    def import_tree(
+        self, tree: ExportMessageTree, state: TreeState = TreeState.BACKLOG_RANKING
+    ) -> tuple[MessageTreeState, Message]:
+        assert tree.message_tree_id is not None and tree.message_tree_id == tree.prompt.message_id
+        root_msg = self.import_message(tree.prompt, message_tree_id=tree.prompt.message_id)
+        assert state == TreeState.BACKLOG_RANKING or state == TreeState.RANKING, f"{state} not supported for import"
+        active = state == TreeState.RANKING
+        mts = MessageTreeState(
+            message_tree_id=root_msg.id,
+            goal_tree_size=0,
+            max_depth=0,
+            max_children_count=0,
+            state=state,
+            origin=self.origin,
+            active=active,
+        )
+        self.db.add(mts)
+        return mts, root_msg
+
+
+def import_file(
+    input_file_path: Path,
+    origin: str,
+    *,
+    model_name: Optional[str] = None,
+    num_activate: int = 0,
+    max_count: Optional[int] = None,
+    dry_run: bool = False,
+) -> int:
+    @db_utils.managed_tx_function(auto_commit=db_utils.CommitMode.ROLLBACK if dry_run else db_utils.CommitMode.COMMIT)
+    def import_tx(db: Session) -> int:
+        importer = Importer(db, origin=origin, model_name=model_name)
+        i = 0
+        with input_file_path.open() as file_in:
+            # read line tree object
+            for line in file_in:
+                dict_tree = json.loads(line)
+
+                # validate data
+                tree: ExportMessageTree = pydantic.parse_obj_as(ExportMessageTree, dict_tree)
+                existing_mts = importer.fetch_message_tree_state(tree.message_tree_id)
+                if existing_mts:
+                    logger.info(f"Skipping existing message tree: {tree.message_tree_id}")
+                else:
+                    state = TreeState.BACKLOG_RANKING if i >= num_activate else TreeState.RANKING
+                    mts, root_msg = importer.import_tree(tree, state=state)
+                    i += 1
+                    logger.info(
+                        f"imported tree: {mts.message_tree_id}, {mts.state=}, {mts.active=}, {root_msg.children_count=}"
+                    )
+
+                if max_count and i >= max_count:
+                    logger.info(f"Reached max count {max_count} of trees to import.")
+                    break
+        return i
+
+    if dry_run:
+        logger.info("DRY RUN with rollback")
+    return import_tx()
+
+
+def parse_args():
+    def str2bool(v):
+        if isinstance(v, bool):
+            return v
+        if v.lower() in ("yes", "true", "t", "y", "1"):
+            return True
+        elif v.lower() in ("no", "false", "f", "n", "0"):
+            return False
+        else:
+            raise argparse.ArgumentTypeError("Boolean value expected.")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "input_file_path",
+        help="Input file path",
+    )
+    parser.add_argument("--origin", type=str, default=None, help="Value for origin of message trees")
+    parser.add_argument("--model_name", type=str, default=None, help="Default name of model (if missing in messages)")
+    parser.add_argument("--num_activate", type=int, default=0, help="Number of trees to add in ranking state")
+    parser.add_argument("--max_count", type=int, default=None, help="Maximum number of message trees to import")
+    parser.add_argument("--dry_run", type=str2bool, default=False)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    input_file_path = Path(args.input_file_path)
+    if not input_file_path.exists() or not input_file_path.is_file():
+        print("Invalid input file:", args.input_file_path)
+        exit(1)
+
+    dry_run = args.dry_run
+    num_imported = import_file(
+        input_file_path,
+        origin=args.origin or input_file_path.name,
+        model_name=args.model_name,
+        num_activate=args.num_activate,
+        max_count=args.max_count,
+        dry_run=dry_run,
+    )
+
+    logger.info(f"Done ({num_imported=}, {dry_run=})")
+
+
+if __name__ == "__main__":
+    main()

backend/main.py

@@ -191,6 +191,7 @@ class DummyMessage(BaseModel):
                         review_count=5,
                         review_result=True,
                         check_tree_state=False,
+                        check_duplicate=False,
                     )
                     if message.parent_id is None:
                         tm._insert_default_state(
@@ -215,7 +216,8 @@ def ensure_tree_states():
     try:
         logger.info("Startup: TreeManager.ensure_tree_states()")
         with Session(engine) as db:
-            tm = TreeManager(db, None)
+            api_client = api_auth(settings.OFFICIAL_WEB_API_KEY, db=db)
+            tm = TreeManager(db, PromptRepository(db, api_client=api_client))
             tm.ensure_tree_states()
 
     except Exception:

backend/oasst_backend/api/deps.py

@@ -1,6 +1,7 @@
 from http import HTTPStatus
 from secrets import token_hex
-from typing import Generator, NamedTuple
+from typing import Generator, NamedTuple, Optional
+from uuid import UUID
 
 from fastapi import Depends, Request, Response, Security
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
@@ -67,6 +68,7 @@ def create_api_client(
     trusted: bool | None = False,
     admin_email: str | None = None,
     api_key: str | None = None,
+    force_id: Optional[UUID] = None,
 ) -> ApiClient:
     if api_key is None:
         api_key = token_hex(32)
@@ -79,6 +81,8 @@ def create_api_client(
         trusted=trusted,
         admin_email=admin_email,
     )
+    if force_id:
+        api_client.id = force_id
     session.add(api_client)
     session.commit()
     session.refresh(api_client)

backend/oasst_backend/config.py

@@ -46,6 +46,15 @@ class TreeManagerConfiguration(BaseModel):
     num_required_rankings: int = 3
     """Number of rankings in which the message participated."""
 
+    p_activate_backlog_tree: float = 0.8
+    """Probability to activate a message tree in BACKLOG_RANKING state when another tree enters
+    a terminal state. Use this settting to control ratio of initial prompts and backlog tree
+    activations."""
+
+    min_active_rankings_per_lang: int = 2
+    """When the number of active ranking tasks is below this value when a tree enters a terminal
+    state an available trees in BACKLOG_RANKING will be actived (i.e. enters the RANKING state)."""
+
     labels_initial_prompt: list[TextLabel] = [
         TextLabel.spam,
         TextLabel.quality,

backend/oasst_backend/models/message.py

@@ -57,6 +57,11 @@ def __new__(cls, *args: Any, **kwargs: Any):
 
     rank: Optional[int] = Field(nullable=True)
 
+    synthetic: Optional[bool] = Field(
+        sa_column=sa.Column(sa.Boolean, default=False, server_default=false(), nullable=False)
+    )
+    model_name: Optional[str] = Field(sa_column=sa.Column(sa.String(1024), nullable=True))
+
     emojis: Optional[dict[str, int]] = Field(default=None, sa_column=sa.Column(pg.JSONB), nullable=False)
     _user_emojis: Optional[list[str]] = PrivateAttr(default=None)
 

backend/oasst_backend/models/message_tree_state.py

@@ -43,6 +43,9 @@ class State(str, Enum):
     HALTED_BY_MODERATOR = "halted_by_moderator"
     """A moderator decided to manually halt the message tree construction process."""
 
+    BACKLOG_RANKING = "backlog_ranking"
+    """Imported tree ready to be activated and ranked by users (currently inactive)."""
+
 
 VALID_STATES = (
     State.INITIAL_PROMPT_REVIEW,
@@ -51,6 +54,7 @@ class State(str, Enum):
     State.READY_FOR_SCORING,
     State.READY_FOR_EXPORT,
     State.ABORTED_LOW_GRADE,
+    State.BACKLOG_RANKING,
 )
 
 TERMINAL_STATES = (State.READY_FOR_EXPORT, State.ABORTED_LOW_GRADE, State.SCORING_FAILED, State.HALTED_BY_MODERATOR)
@@ -67,3 +71,4 @@ class MessageTreeState(SQLModel, table=True):
     max_children_count: int = Field(nullable=False)
     state: str = Field(nullable=False, max_length=128, index=True)
     active: bool = Field(nullable=False, index=True)
+    origin: str = Field(sa_column=sa.Column(sa.String(1024), nullable=True))

backend/oasst_backend/prompt_repository.py

@@ -177,6 +177,7 @@ def store_text_reply(
         review_count: int = 0,
         review_result: bool = False,
         check_tree_state: bool = True,
+        check_duplicate: bool = True,
     ) -> Message:
         self.ensure_user_is_enabled()
 
@@ -199,7 +200,7 @@ def store_text_reply(
             logger.error(f"Message size {len(text)=} exceeds size limit of {settings.MESSAGE_SIZE_LIMIT=}.")
             raise OasstError("Message size too long.", OasstErrorCode.TASK_MESSAGE_TOO_LONG)
 
-        if self.check_users_recent_replies_for_duplicates(text):
+        if check_duplicate and self.check_users_recent_replies_for_duplicates(text):
             raise OasstError("User recent messages have duplicates", OasstErrorCode.TASK_MESSAGE_DUPLICATED)
 
         if task.parent_message_id:
@@ -909,8 +910,7 @@ def update_children_counts(self, message_tree_id: UUID):
 ) AS cc
 WHERE message.id = cc.id;
 """
-        r = self.db.execute(text(sql_update_children_count), {"message_tree_id": message_tree_id})
-        logger.debug(f"update_children_count({message_tree_id=}): {r.rowcount} rows.")
+        self.db.execute(text(sql_update_children_count), {"message_tree_id": message_tree_id})
 
     @managed_tx_method(CommitMode.COMMIT)
     def mark_messages_deleted(self, messages: Message | UUID | list[Message | UUID], recursive: bool = True):