From 1153734edb1db9c60693d552025cd92f18c9f7d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20K=C3=B6pf?= Date: Tue, 7 Feb 2023 22:17:36 +0100 Subject: [PATCH] Allow to export filtered by tree state and/or only prompts (#1314) * allow to export filtered by tree state and/or only prompts * formatting * add option to filter by lang code * remove old export code * export as list in only-deleted case --- backend/export.py | 146 ++++++++++++++------- backend/main.py | 35 +---- backend/oasst_backend/prompt_repository.py | 6 - backend/oasst_backend/tree_manager.py | 44 ------- backend/oasst_backend/utils/tree_export.py | 62 +++++++-- 5 files changed, 151 insertions(+), 142 deletions(-) diff --git a/backend/export.py b/backend/export.py index ef3fde49..c0164172 100644 --- a/backend/export.py +++ b/backend/export.py @@ -1,36 +1,43 @@ import argparse -import json -import sys from pathlib import Path from typing import List, Optional from uuid import UUID -from fastapi.encoders import jsonable_encoder from loguru import logger from oasst_backend.database import engine -from oasst_backend.models import Message, MessageTreeState, message_tree_state +from oasst_backend.models import Message, MessageTreeState +from oasst_backend.models.message_tree_state import State as TreeState from oasst_backend.utils import tree_export -from sqlmodel import Session, not_ +from sqlmodel import Session -def fetch_tree_ids(db: Session, ready_only: bool = False): - tree_qry = db.query(MessageTreeState) +def fetch_tree_ids( + db: Session, + state_filter: Optional[TreeState] = None, + lang: Optional[str] = None, +) -> list[tuple[UUID, TreeState]]: + tree_qry = ( + db.query(MessageTreeState) + .select_from(MessageTreeState) + .join(Message, MessageTreeState.message_tree_id == Message.id) + ) - if ready_only: - tree_qry = tree_qry.filter(MessageTreeState.state == message_tree_state.State.READY_FOR_EXPORT) + if lang is not None: + tree_qry = tree_qry.filter(Message.lang == lang) - message_trees = tree_qry.all() + if state_filter: + tree_qry = tree_qry.filter(MessageTreeState.state == state_filter) - message_tree_ids = [tree.message_tree_id for tree in message_trees] - return message_tree_ids + return [(tree.message_tree_id, tree.state) for tree in tree_qry] -def fetch_message_ids( +def fetch_tree_messages( db: Session, message_tree_id: Optional[UUID] = None, user_id: Optional[UUID] = None, - include_deleted: bool = False, - deleted_only: bool = False, + deleted: bool = None, + prompts_only: bool = False, + lang: Optional[str] = None, ) -> List[Message]: qry = db.query(Message) @@ -38,10 +45,12 @@ def fetch_message_ids( qry = qry.filter(Message.message_tree_id == message_tree_id) if user_id: qry = qry.filter(Message.user_id == user_id) - if not include_deleted: - qry = qry.filter(not_(Message.deleted)) - if deleted_only: - qry = qry.filter(Message.deleted) + if deleted is not None: + qry = qry.filter(Message.deleted == deleted) + if prompts_only: + qry = qry.filter(Message.parent_id.is_(None)) + if lang: + qry = qry.filter(Message.lang == lang) return qry.all() @@ -50,31 +59,49 @@ def export_trees( db: Session, export_file: Optional[Path] = None, use_compression: bool = False, - ready_only: bool = False, - include_deleted: bool = False, - deleted_only: bool = False, + deleted: bool = False, user_id: Optional[UUID] = None, -): + prompts_only: bool = False, + state_filter: Optional[TreeState] = None, + lang: Optional[str] = None, +) -> None: trees_to_export: List[tree_export.ExportMessageTree] = [] if user_id: - messages = fetch_message_ids(db, user_id=user_id, include_deleted=include_deleted, deleted_only=deleted_only) - message_tree_ids = [msg.message_tree_id for msg in messages] + messages = fetch_tree_messages( + db, + user_id=user_id, + deleted=deleted, + prompts_only=prompts_only, + lang=lang, + ) + tree_export.write_messages_to_file(export_file, messages, use_compression) else: - message_tree_ids = fetch_tree_ids(db, ready_only) + message_tree_ids = fetch_tree_ids(db, state_filter, lang=lang) - message_trees = [ - fetch_message_ids(db, message_tree_id=tree_id, include_deleted=include_deleted, deleted_only=deleted_only) - for tree_id in message_tree_ids - ] + message_trees = [ + fetch_tree_messages( + db, + message_tree_id=tree_id, + deleted=deleted, + prompts_only=prompts_only, + lang=None, + ) + for (tree_id, _) in message_tree_ids + ] - for message_tree_id, message_tree in zip(message_tree_ids, message_trees): - trees_to_export.append(tree_export.build_export_tree(message_tree_id, message_tree)) + # when exporting only-deleted we don't have a porper tree structure, export as list + if deleted is True: + messages = [m for t in message_trees for m in t] + tree_export.write_messages_to_file(export_file, messages, use_compression) + else: + for (message_tree_id, message_tree_state), message_tree in zip(message_tree_ids, message_trees): + t = tree_export.build_export_tree(message_tree_id, message_tree_state, message_tree) + if prompts_only: + t.prompt.replies = None + trees_to_export.append(t) - if export_file: - tree_export.write_trees_to_file(export_file, trees_to_export, use_compression) - else: - sys.stdout.write(json.dumps(jsonable_encoder(trees_to_export), indent=2)) + tree_export.write_trees_to_file(export_file, trees_to_export, use_compression) def validate_args(args): @@ -83,8 +110,8 @@ def validate_args(args): args.use_compression = args.export_file is not None and ".gz" in args.export_file - if args.ready_only and args.user_id is not None: - raise ValueError("Cannot use --ready-only when specifying a user ID") + if args.state and args.user is not None: + raise ValueError("Cannot use --state when specifying a user ID") if args.export_file is None: logger.warning("No export file provided, output will be sent to STDOUT") @@ -98,11 +125,6 @@ def parse_args(): type=str, help="Name of file to export trees to. If not provided, output will be sent to STDOUT", ) - parser.add_argument( - "--ready-only", - action="store_true", - help="Only export trees which are ready for use", - ) parser.add_argument( "--include-deleted", action="store_true", @@ -116,7 +138,22 @@ def parse_args(): parser.add_argument( "--user", type=str, - help="Only export trees involving the user with the specified ID. Incompatible with --ready-only.", + help="Only export trees involving the user with the specified ID. Incompatible with --state.", + ) + parser.add_argument( + "--state", + type=str, + help="all|prompt_lottery_waiting|growing|ready_for_export|aborted_low_grade|halted_by_moderator|backlog_ranking", + ) + parser.add_argument( + "--lang", + type=str, + help="Filter message trees by language code (BCP 47)", + ) + parser.add_argument( + "--prompts-only", + action="store_true", + help="Export a list of initial prompt messages", ) args = parser.parse_args() @@ -127,15 +164,28 @@ def main(): args = parse_args() validate_args(args) + state_filter: Optional[TreeState] = None + if args.state is None: + state_filter = TreeState.READY_FOR_EXPORT + elif args.state != "all": + state_filter = TreeState(args.state) + + deleted: Optional[bool] = False + if args.include_deleted: + deleted = None + if args.deleted_only: + deleted = True + with Session(engine) as db: export_trees( db, Path(args.export_file) if args.export_file is not None else None, args.use_compression, - args.ready_only, - args.include_deleted, - args.deleted_only, - UUID(args.user) if args.user is not None else None, + deleted=deleted, + user_id=UUID(args.user) if args.user is not None else None, + prompts_only=args.prompts_only, + state_filter=state_filter, + lang=args.lang, ) diff --git a/backend/main.py b/backend/main.py index dacd8097..807d52ce 100644 --- a/backend/main.py +++ b/backend/main.py @@ -333,24 +333,6 @@ def get_openapi_schema(): return json.dumps(app.openapi()) -def export_ready_trees(file: Optional[str] = None, use_compression: bool = False): - try: - with Session(engine) as db: - api_client = api_auth(settings.OFFICIAL_WEB_API_KEY, db=db) - dummy_user = protocol_schema.User(id="__dummy_user__", display_name="Dummy User", auth_method="local") - - ur = UserRepository(db=db, api_client=api_client) - tr = TaskRepository(db=db, api_client=api_client, client_user=dummy_user, user_repository=ur) - pr = PromptRepository( - db=db, api_client=api_client, client_user=dummy_user, user_repository=ur, task_repository=tr - ) - tm = TreeManager(db, pr) - - tm.export_all_ready_trees(file, use_compression=use_compression) - except Exception: - logger.exception("Error exporting trees.") - - def retry_scoring_failed_message_trees(): try: logger.info("TreeManager.retry_scoring_failed_message_trees()") @@ -382,17 +364,6 @@ def main(): ) parser.add_argument("--host", help="The host to run the server", default="0.0.0.0") parser.add_argument("--port", help="The port to run the server", default=8080) - parser.add_argument( - "--export", - default=False, - help="Export all trees which are ready for exporting.", - action="store_true", - ) - parser.add_argument( - "--export-file", - type=str, - help="Name of file to export trees to. If not provided when exporting, output will be send to STDOUT", - ) parser.add_argument( "--retry-scoring", default=False, @@ -405,14 +376,10 @@ def main(): if args.print_openapi_schema: print(get_openapi_schema()) - if args.export: - use_compression: bool = ".gz" in args.export_file - export_ready_trees(file=args.export_file, use_compression=use_compression) - if args.retry_scoring: retry_scoring_failed_message_trees() - if not (args.export or args.print_openapi_schema or args.retry_scoring): + if not (args.print_openapi_schema or args.retry_scoring): uvicorn.run(app, host=args.host, port=args.port) diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 71e7ab68..c5fb0208 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -656,12 +656,6 @@ class PromptRepository: qry = qry.filter(not_(Message.deleted)) return self._add_user_emojis_all(qry) - def fetch_message_trees_ready_for_export(self) -> list[MessageTreeState]: - qry = self.db.query(MessageTreeState).filter( - MessageTreeState.state == message_tree_state.State.READY_FOR_EXPORT - ) - return qry.all() - def fetch_multiple_random_replies(self, max_size: int = 5, message_role: str = None): """ Fetch a conversation with multiple possible replies to it. diff --git a/backend/oasst_backend/tree_manager.py b/backend/oasst_backend/tree_manager.py index bb541996..51eedac8 100644 --- a/backend/oasst_backend/tree_manager.py +++ b/backend/oasst_backend/tree_manager.py @@ -1,6 +1,4 @@ -import json import random -import sys from datetime import datetime, timedelta from enum import Enum from http import HTTPStatus @@ -10,7 +8,6 @@ from uuid import UUID import numpy as np import pydantic import sqlalchemy as sa -from fastapi.encoders import jsonable_encoder from loguru import logger from oasst_backend.api.v1.utils import prepare_conversation, prepare_conversation_message_list from oasst_backend.config import TreeManagerConfiguration, settings @@ -25,7 +22,6 @@ from oasst_backend.models import ( message_tree_state, ) from oasst_backend.prompt_repository import PromptRepository -from oasst_backend.utils import tree_export from oasst_backend.utils.database_utils import ( CommitMode, async_managed_tx_method, @@ -1673,44 +1669,6 @@ DELETE FROM user_stats WHERE user_id = :user_id; if ban: self.db.execute(update(User).filter(User.id == user_id).values(deleted=True, enabled=False)) - def export_trees_to_file( - self, - message_tree_ids: list[str], - file=None, - reviewed: bool = True, - include_deleted: bool = False, - use_compression: bool = False, - ) -> None: - trees_to_export: List[tree_export.ExportMessageTree] = [] - - for message_tree_id in message_tree_ids: - messages: List[Message] = self.pr.fetch_message_tree(message_tree_id, reviewed, include_deleted) - trees_to_export.append(tree_export.build_export_tree(message_tree_id, messages)) - - if file: - tree_export.write_trees_to_file(file, trees_to_export, use_compression) - else: - sys.stdout.write(json.dumps(jsonable_encoder(trees_to_export), indent=2)) - - def export_all_ready_trees( - self, file: str, reviewed: bool = True, include_deleted: bool = False, use_compression: bool = False - ) -> None: - message_tree_states: MessageTreeState = self.pr.fetch_message_trees_ready_for_export() - message_tree_ids = [ms.message_tree_id for ms in message_tree_states] - self.export_trees_to_file(message_tree_ids, file, reviewed, include_deleted, use_compression) - - def export_all_user_trees( - self, - user_id: str, - file: str, - reviewed: bool = True, - include_deleted: bool = False, - use_compression: bool = False, - ) -> None: - messages = self.pr.fetch_user_message_trees(UUID(user_id)) - message_tree_ids = [ms.message_tree_id for ms in messages] - self.export_trees_to_file(message_tree_ids, file, reviewed, include_deleted, use_compression) - @managed_tx_method(CommitMode.COMMIT) def retry_scoring_failed_message_trees(self): query = self.db.query(MessageTreeState).filter( @@ -1777,5 +1735,3 @@ if __name__ == "__main__": # print( # ".query_tree_ranking_results", tm.query_tree_ranking_results(UUID("21f9d585-d22c-44ab-a696-baa3d83b5f1b")) # ) - - # print(tm.export_trees_to_file(message_tree_ids=["7e75fb38-e664-4e2b-817c-b9a0b01b0074"], file="lol.jsonl")) diff --git a/backend/oasst_backend/utils/tree_export.py b/backend/oasst_backend/utils/tree_export.py index e6bea0d9..409d8a22 100644 --- a/backend/oasst_backend/utils/tree_export.py +++ b/backend/oasst_backend/utils/tree_export.py @@ -1,12 +1,16 @@ from __future__ import annotations +import contextlib import gzip import json +import sys from collections import defaultdict -from typing import Optional, TextIO +from typing import Iterable, Optional, TextIO +from uuid import UUID from fastapi.encoders import jsonable_encoder from oasst_backend.models import Message +from oasst_backend.models.message_tree_state import State as TreeState from pydantic import BaseModel @@ -41,10 +45,13 @@ class ExportMessageNode(BaseModel): class ExportMessageTree(BaseModel): message_tree_id: str + tree_state: Optional[str] prompt: Optional[ExportMessageNode] -def build_export_tree(message_tree_id: str, messages: list[Message]) -> ExportMessageTree: +def build_export_tree( + message_tree_id: UUID, message_tree_state: TreeState, messages: list[Message] +) -> ExportMessageTree: export_messages = [ExportMessageNode.prep_message_export(m) for m in messages] messages_by_parent = defaultdict(list) @@ -59,19 +66,54 @@ def build_export_tree(message_tree_id: str, messages: list[Message]) -> ExportMe return node prompt = assign_replies(messages_by_parent[None][0]) - return ExportMessageTree(message_tree_id=str(message_tree_id), prompt=prompt) + return ExportMessageTree(message_tree_id=str(message_tree_id), tree_state=message_tree_state, prompt=prompt) -def write_trees_to_file(file, trees: list[ExportMessageTree], use_compression: bool = True) -> None: - - out_buff: TextIO - if use_compression: - out_buff = gzip.open(file, "wt", encoding="UTF-8") +# see https://stackoverflow.com/questions/17602878/how-to-handle-both-with-open-and-sys-stdout-nicely +@contextlib.contextmanager +def smart_open(filename: str = None) -> TextIO: + if filename and filename != "-": + fh = open(filename, "wt", encoding="UTF-8") else: - out_buff = open(file, "wt", encoding="UTF-8") + fh = sys.stdout + + try: + yield fh + finally: + if fh is not sys.stdout: + fh.close() + + +def write_trees_to_file(filename: str | None, trees: list[ExportMessageTree], use_compression: bool = True) -> None: + out_buff: TextIO + + if use_compression: + if not filename: + raise RuntimeError("File name must be specified when using compression.") + out_buff = gzip.open(filename, "wt", encoding="UTF-8") + else: + out_buff = smart_open(filename) with out_buff as f: for tree in trees: - file_data = jsonable_encoder(tree) + file_data = jsonable_encoder(tree, exclude_none=True) + json.dump(file_data, f) + f.write("\n") + + +def write_messages_to_file(filename: str | None, messages: Iterable[Message], use_compression: bool = True) -> None: + out_buff: TextIO + + if use_compression: + if not filename: + raise RuntimeError("File name must be specified when using compression.") + out_buff = gzip.open(filename, "wt", encoding="UTF-8") + else: + out_buff = smart_open(filename) + + with out_buff as f: + for m in messages: + export_message = ExportMessageNode.prep_message_export(m) + file_data = jsonable_encoder(export_message, exclude_none=True) json.dump(file_data, f) f.write("\n")