From e0df9f0b7c70da404c8c1049737c00027f0c2a0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20K=C3=B6pf?= Date: Wed, 1 Feb 2023 23:49:21 +0100 Subject: [PATCH] Add infty prompting + prompt-lottery (#1055) * first bits of prompt-lottery * add tree prompt lottery * minor text changes --- ...8ff8_add_won_prompt_lottery_date_to_mts.py | 27 +++++ backend/oasst_backend/config.py | 4 +- .../models/message_tree_state.py | 7 ++ backend/oasst_backend/tree_manager.py | 105 +++++++++++++++--- 4 files changed, 126 insertions(+), 17 deletions(-) create mode 100644 backend/alembic/versions/2023_02_01_1010-f60958968ff8_add_won_prompt_lottery_date_to_mts.py diff --git a/backend/alembic/versions/2023_02_01_1010-f60958968ff8_add_won_prompt_lottery_date_to_mts.py b/backend/alembic/versions/2023_02_01_1010-f60958968ff8_add_won_prompt_lottery_date_to_mts.py new file mode 100644 index 00000000..f82f28fd --- /dev/null +++ b/backend/alembic/versions/2023_02_01_1010-f60958968ff8_add_won_prompt_lottery_date_to_mts.py @@ -0,0 +1,27 @@ +"""add won_prompt_lottery_date to mts + +Revision ID: f60958968ff8 +Revises: 7b8f0011e0b0 +Create Date: 2023-02-01 10:10:38.301707 + +""" +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "f60958968ff8" +down_revision = "7b8f0011e0b0" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("message_tree_state", sa.Column("won_prompt_lottery_date", sa.DateTime(timezone=True), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("message_tree_state", "won_prompt_lottery_date") + # ### end Alembic commands ### diff --git a/backend/oasst_backend/config.py b/backend/oasst_backend/config.py index 449543d5..17935021 100644 --- a/backend/oasst_backend/config.py +++ b/backend/oasst_backend/config.py @@ -13,7 +13,7 @@ class TreeManagerConfiguration(BaseModel): No new initial prompt tasks are handed out to users if this number is reached.""" - max_tree_depth: int = 6 + max_tree_depth: int = 3 """Maximum depth of message tree.""" max_children_count: int = 3 @@ -22,7 +22,7 @@ class TreeManagerConfiguration(BaseModel): num_prompter_replies: int = 1 """Number of prompter replies to collect per assistant reply.""" - goal_tree_size: int = 15 + goal_tree_size: int = 12 """Total number of messages to gather per tree.""" num_reviews_initial_prompt: int = 3 diff --git a/backend/oasst_backend/models/message_tree_state.py b/backend/oasst_backend/models/message_tree_state.py index a286d483..199f475b 100644 --- a/backend/oasst_backend/models/message_tree_state.py +++ b/backend/oasst_backend/models/message_tree_state.py @@ -1,4 +1,6 @@ +from datetime import datetime from enum import Enum +from typing import Optional from uuid import UUID import sqlalchemy as sa @@ -46,6 +48,9 @@ class State(str, Enum): BACKLOG_RANKING = "backlog_ranking" """Imported tree ready to be activated and ranked by users (currently inactive).""" + PROMPT_LOTTERY_WAITING = "prompt_lottery_waiting" + """Initial prompt has passed spam check, waiting to be drawn to grow.""" + VALID_STATES = ( State.INITIAL_PROMPT_REVIEW, @@ -63,6 +68,7 @@ TERMINAL_STATES = ( State.SCORING_FAILED, State.HALTED_BY_MODERATOR, State.BACKLOG_RANKING, + State.PROMPT_LOTTERY_WAITING, ) @@ -78,3 +84,4 @@ class MessageTreeState(SQLModel, table=True): 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)) + won_prompt_lottery_date: Optional[datetime] = Field(sa_column=sa.Column(sa.DateTime(timezone=True), nullable=True)) diff --git a/backend/oasst_backend/tree_manager.py b/backend/oasst_backend/tree_manager.py index b1d50d35..709fc622 100644 --- a/backend/oasst_backend/tree_manager.py +++ b/backend/oasst_backend/tree_manager.py @@ -21,6 +21,7 @@ from oasst_backend.utils.hugging_face import HfClassificationModel, HfEmbeddingM from oasst_backend.utils.ranking import ranked_pairs from oasst_shared.exceptions.oasst_api_error import OasstError, OasstErrorCode from oasst_shared.schemas import protocol as protocol_schema +from oasst_shared.utils import utcnow from sqlmodel import Session, func, not_, or_, text, update @@ -153,7 +154,7 @@ class TreeManager: def _determine_task_availability_internal( self, - num_active_trees: int, + num_missing_prompts: int, extendible_parents: list[ExtendibleParentRow], prompts_need_review: list[Message], replies_need_review: list[Message], @@ -161,8 +162,7 @@ class TreeManager: ) -> dict[protocol_schema.TaskRequestType, int]: task_count_by_type: dict[protocol_schema.TaskRequestType, int] = {t: 0 for t in protocol_schema.TaskRequestType} - num_missing_prompts = max(0, self.cfg.max_active_trees - num_active_trees) - task_count_by_type[protocol_schema.TaskRequestType.initial_prompt] = num_missing_prompts + task_count_by_type[protocol_schema.TaskRequestType.initial_prompt] = max(1, num_missing_prompts) task_count_by_type[protocol_schema.TaskRequestType.prompter_reply] = len( list(filter(lambda x: x.parent_role == "assistant", extendible_parents)) @@ -194,6 +194,72 @@ class TreeManager: return task_count_by_type + def _prompt_lottery(self, lang: str) -> int: + MAX_RETRIES = 5 + retry = 0 + while True: + num_active_trees = self.query_num_active_trees(lang=lang, exclude_ranking=True) + num_missing_prompts = self.cfg.max_active_trees - num_active_trees + if num_missing_prompts <= 0: + return 0 + + # select among distinct users + authors_qry = ( + self.db.query(Message.user_id) + .select_from(MessageTreeState) + .join(Message, MessageTreeState.message_tree_id == Message.id) + .filter( + MessageTreeState.state == message_tree_state.State.PROMPT_LOTTERY_WAITING, + Message.lang == lang, + not_(Message.deleted), + Message.review_result, + ) + .distinct(Message.user_id) + ) + + author_ids = authors_qry.all() + if len(author_ids) == 0: + logger.info( + f"No prompts for prompt lottery available ({num_missing_prompts} trees missing for {lang=})." + ) + return num_missing_prompts + + # first select an authour + prompt_author_id: UUID = random.choice(author_ids)["user_id"] + logger.info(f"Selected random prompt author {prompt_author_id} among {len(author_ids)} candidates.") + + # select random prompt of author + qry = ( + self.db.query(MessageTreeState, Message) + .select_from(MessageTreeState) + .join(Message, MessageTreeState.message_tree_id == Message.id) + .filter( + MessageTreeState.state == message_tree_state.State.PROMPT_LOTTERY_WAITING, + Message.user_id == prompt_author_id, + Message.lang == lang, + not_(Message.deleted), + Message.review_result, + ) + .limit(100) + ) + + prompt_candidates = qry.all() + if len(prompt_candidates) == 0: + retry += 1 # not sure if this can happen with repeatable read isolation level, just in case we retry + if retry < MAX_RETRIES: + continue + else: + logger.warning("Max retries in prompt lottery reached.") + return num_missing_prompts + + winner_prompt = random.choice(prompt_candidates) + message: Message = winner_prompt.Message + logger.info(f"Prompt lottery winner: {message.id=}") + + mts: MessageTreeState = winner_prompt.MessageTreeState + self._enter_state(mts, message_tree_state.State.GROWING) + self.db.flush() + def determine_task_availability(self, lang: str) -> dict[protocol_schema.TaskRequestType, int]: self.pr.ensure_user_is_enabled() @@ -201,14 +267,14 @@ class TreeManager: lang = "en" logger.warning("Task availability request without lang tag received, assuming lang='en'.") - num_active_trees = self.query_num_active_trees(lang=lang, exclude_ranking=True) + num_missing_prompts = self._prompt_lottery(lang=lang) extendible_parents, _ = self.query_extendible_parents(lang=lang) prompts_need_review = self.query_prompts_need_review(lang=lang) replies_need_review = self.query_replies_need_review(lang=lang) incomplete_rankings = self.query_incomplete_rankings(lang=lang) return self._determine_task_availability_internal( - num_active_trees=num_active_trees, + num_missing_prompts=num_missing_prompts, extendible_parents=extendible_parents, prompts_need_review=prompts_need_review, replies_need_review=replies_need_review, @@ -238,7 +304,8 @@ class TreeManager: lang = "en" logger.warning("Task request without lang tag received, assuming 'en'.") - num_active_trees = self.query_num_active_trees(lang=lang, exclude_ranking=True) + num_missing_prompts = self._prompt_lottery(lang=lang) + prompts_need_review = self.query_prompts_need_review(lang=lang) replies_need_review = self.query_replies_need_review(lang=lang) extendible_parents, active_tree_sizes = self.query_extendible_parents(lang=lang) @@ -256,7 +323,7 @@ class TreeManager: num_ranking_tasks=len(incomplete_rankings), num_replies_need_review=len(replies_need_review), num_prompts_need_review=len(prompts_need_review), - num_missing_prompts=max(0, self.cfg.max_active_trees - num_active_trees), + num_missing_prompts=num_missing_prompts, num_missing_replies=num_missing_replies, ) @@ -268,7 +335,7 @@ class TreeManager: ) else: task_count_by_type = self._determine_task_availability_internal( - num_active_trees=num_active_trees, + num_missing_prompts=num_missing_prompts, extendible_parents=extendible_parents, prompts_need_review=prompts_need_review, replies_need_review=replies_need_review, @@ -611,7 +678,7 @@ class TreeManager: ) else: self.enter_low_grade_state(msg.message_tree_id) - self.check_condition_for_growing_state(msg.message_tree_id) + self.check_condition_for_prompt_lottery(msg.message_tree_id) elif msg.review_count >= self.cfg.num_reviews_reply: if not msg.review_result and acceptance_score > self.cfg.acceptance_threshold_reply: msg.review_result = True @@ -649,6 +716,8 @@ class TreeManager: if len(incomplete_rankings) < self.cfg.min_active_rankings_per_lang: self.activate_backlog_tree(lang=root_msg.lang) else: + if mts.state == message_tree_state.State.GROWING and mts.won_prompt_lottery_date is None: + mts.won_prompt_lottery_date = utcnow() logger.info(f"Tree entered '{mts.state}' state ({mts.message_tree_id=})") def enter_low_grade_state(self, message_tree_id: UUID) -> None: @@ -656,8 +725,8 @@ class TreeManager: mts = self.pr.fetch_tree_state(message_tree_id) self._enter_state(mts, message_tree_state.State.ABORTED_LOW_GRADE) - def check_condition_for_growing_state(self, message_tree_id: UUID) -> bool: - logger.debug(f"check_condition_for_growing_state({message_tree_id=})") + def check_condition_for_prompt_lottery(self, message_tree_id: UUID) -> bool: + logger.debug(f"check_condition_for_prompt_lottery({message_tree_id=})") mts = self.pr.fetch_tree_state(message_tree_id) if not mts.active or mts.state != message_tree_state.State.INITIAL_PROMPT_REVIEW: @@ -670,7 +739,7 @@ class TreeManager: logger.debug(f"False {initial_prompt.review_result=}") return False - self._enter_state(mts, message_tree_state.State.GROWING) + self._enter_state(mts, message_tree_state.State.PROMPT_LOTTERY_WAITING) return True def check_condition_for_ranking_state(self, message_tree_id: UUID) -> bool: @@ -1097,7 +1166,7 @@ LEFT JOIN message_reaction mr ON mr.task_id = t.id AND mr.payload_type = 'Rankin f"Checking state of {len(prompt_review_trees)} active message trees in 'initial_prompt_review' state." ) for t in prompt_review_trees: - self.check_condition_for_growing_state(t.message_tree_id) + self.check_condition_for_prompt_lottery(t.message_tree_id) growing_trees: list[MessageTreeState] = ( self.db.query(MessageTreeState) @@ -1288,11 +1357,17 @@ DELETE FROM message WHERE message_tree_id = :message_tree_id; logger.debug(f"purge_message_tree({message_tree_id=}) {r.rowcount} rows.") def _reactivate_tree(self, mts: MessageTreeState): - self._enter_state(mts, message_tree_state.State.INITIAL_PROMPT_REVIEW) + if mts.state == message_tree_state.State.PROMPT_LOTTERY_WAITING: + return + tree_id = mts.message_tree_id - if self.check_condition_for_growing_state(tree_id): + if mts.won_prompt_lottery_date is not None: + self._enter_state(mts, message_tree_state.State.GROWING) if self.check_condition_for_ranking_state(tree_id): self.check_condition_for_scoring_state(tree_id) + else: + self._enter_state(mts, message_tree_state.State.INITIAL_PROMPT_REVIEW) + self.check_condition_for_prompt_lottery(tree_id) @managed_tx_method(CommitMode.FLUSH) def purge_user_messages(