diff --git a/backend/main.py b/backend/main.py index 147f8e30..a9dd8f72 100644 --- a/backend/main.py +++ b/backend/main.py @@ -273,6 +273,24 @@ 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 main(): # Importing here so we don't import packages unnecessarily if we're # importing main as a module. @@ -289,11 +307,21 @@ 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", help="Export all trees which are ready for exporting.", action=argparse.BooleanOptionalAction + ) + parser.add_argument( + "--export-file", + help="Name of file to export trees to. If not provided when exporting, output will be send to STDOUT", + ) args = parser.parse_args() if args.print_openapi_schema: print(get_openapi_schema()) + elif args.export: + use_compression: bool = ".gz" in args.export_file + export_ready_trees(file=args.export_file, use_compression=use_compression) else: 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 8c746320..f9ba9ad6 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -2,7 +2,7 @@ import random from collections import defaultdict from datetime import datetime from http import HTTPStatus -from typing import List, Optional, Tuple +from typing import Optional from uuid import UUID, uuid4 import oasst_backend.models.db_payload as db_payload @@ -255,7 +255,7 @@ class PromptRepository: return reaction @managed_tx_method(CommitMode.COMMIT) - def store_ranking(self, ranking: protocol_schema.MessageRanking) -> Tuple[MessageReaction, Task]: + def store_ranking(self, ranking: protocol_schema.MessageRanking) -> tuple[MessageReaction, Task]: # fetch task task = self.task_repository.fetch_task_by_frontend_message_id(ranking.message_id) self._validate_task(task, frontend_message_id=ranking.message_id) @@ -345,13 +345,13 @@ class PromptRepository: return message_toxicity @managed_tx_method(CommitMode.FLUSH) - def insert_message_embedding(self, message_id: UUID, model: str, embedding: List[float]) -> MessageEmbedding: + def insert_message_embedding(self, message_id: UUID, model: str, embedding: list[float]) -> MessageEmbedding: """Insert the embedding of a new message in the database. Args: message_id (UUID): the identifier of the message we want to save its embedding model (str): the model used for creating the embedding - embedding (List[float]): the values obtained from the message & model + embedding (list[float]): the values obtained from the message & model Raises: OasstError: if misses some of the before params @@ -383,7 +383,7 @@ class PromptRepository: return reaction @managed_tx_method(CommitMode.FLUSH) - def store_text_labels(self, text_labels: protocol_schema.TextLabels) -> Tuple[TextLabels, Task, Message]: + def store_text_labels(self, text_labels: protocol_schema.TextLabels) -> tuple[TextLabels, Task, Message]: valid_labels: Optional[list[str]] = None mandatory_labels: Optional[list[str]] = None @@ -529,6 +529,22 @@ class PromptRepository: qry = qry.filter(not_(Message.deleted)) return qry.all() + def fetch_user_message_trees( + self, user_id: Message.user_id, reviewed: bool = True, include_deleted: bool = False + ) -> list[Message]: + qry = self.db.query(Message).filter(Message.user_id == user_id) + if reviewed: + qry = qry.filter(Message.review_result) + if not include_deleted: + qry = qry.filter(not_(Message.deleted)) + return qry.all() + + 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 d54df08b..48dea1c9 100644 --- a/backend/oasst_backend/tree_manager.py +++ b/backend/oasst_backend/tree_manager.py @@ -1,4 +1,6 @@ +import json import random +import sys from datetime import datetime from enum import Enum from http import HTTPStatus @@ -7,11 +9,13 @@ from uuid import UUID import numpy as np import pydantic +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 from oasst_backend.models import Message, MessageReaction, MessageTreeState, Task, TextLabels, User, 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, managed_tx_method from oasst_backend.utils.hugging_face import HfClassificationModel, HfEmbeddingModel, HfUrl, HuggingFaceAPI from oasst_backend.utils.ranking import ranked_pairs @@ -1184,14 +1188,55 @@ 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) + if __name__ == "__main__": from oasst_backend.api.deps import api_auth + + # from oasst_backend.api.deps import create_api_client from oasst_backend.database import engine from oasst_backend.prompt_repository import PromptRepository with Session(engine) as db: api_client = api_auth(settings.OFFICIAL_WEB_API_KEY, db=db) + # api_client = create_api_client(session=db, description="test", frontend_type="bot") dummy_user = protocol_schema.User(id="__dummy_user__", display_name="Dummy User", auth_method="local") pr = PromptRepository(db=db, api_client=api_client, client_user=dummy_user) @@ -1200,25 +1245,22 @@ if __name__ == "__main__": tm = TreeManager(db, pr, cfg) tm.ensure_tree_states() - tm.purge_user_messages(user_id=UUID("2ef9ad21-0dc5-442d-8750-6f7f1790723f"), purge_initial_prompts=False) + # tm.purge_user_messages(user_id=UUID("2ef9ad21-0dc5-442d-8750-6f7f1790723f"), purge_initial_prompts=False) # tm.purge_user(user_id=UUID("2ef9ad21-0dc5-442d-8750-6f7f1790723f")) # db.commit() # print("query_num_active_trees", tm.query_num_active_trees()) # print("query_incomplete_rankings", tm.query_incomplete_rankings()) # print("query_replies_need_review", tm.query_replies_need_review()) + # print("query_incomplete_reply_reviews", tm.query_replies_need_review()) # print("query_incomplete_initial_prompt_reviews", tm.query_prompts_need_review()) # print("query_extendible_trees", tm.query_extendible_trees()) # print("query_extendible_parents", tm.query_extendible_parents()) - # print("query_tree_size", tm.query_tree_size(message_tree_id=UUID("bdf434cf-4df5-4b74-949c-a5a157bc3292"))) - - # print( - # "query_reviews_for_message", - # tm.query_reviews_for_message(message_id=UUID("6a444493-0d48-4316-a9f1-7e263f5a2473")), - # ) # print("next_task:", tm.next_task()) # print( - # "query_tree_ranking_results", tm.query_tree_ranking_results(UUID("6036f58f-41b5-48c4-bdd9-b16f34ab1312")) + # ".query_tree_ranking_results", tm.query_tree_ranking_results(UUID("2ac20d38-6650-43aa-8bb3-f61080c0d921")) # ) + + 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 new file mode 100644 index 00000000..ee3de9d7 --- /dev/null +++ b/backend/oasst_backend/utils/tree_export.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import gzip +import json +from collections import defaultdict +from typing import Optional, TextIO + +from fastapi.encoders import jsonable_encoder +from oasst_backend.models import Message +from pydantic import BaseModel + + +class ExportMessageNode(BaseModel): + message_id: str + parent_id: Optional[str] + text: Optional[str] + role: str + review_count: Optional[int] + rank: Optional[int] + replies: Optional[list[ExportMessageNode]] + + @classmethod + def prep_message_export(cls, message: Message) -> ExportMessageNode: + return cls( + message_id=str(message.id), + parent_id=str(message.parent_id) if message.parent_id else None, + text=str(message.payload.payload.text), + role=message.role, + review_count=message.review_count, + rank=message.rank, + ) + + +class ExportMessageTree(BaseModel): + message_tree_id: str + replies: Optional[ExportMessageNode] + + +def build_export_tree(message_tree_id: str, messages: list[Message]) -> ExportMessageTree: + export_tree = ExportMessageTree(message_tree_id=str(message_tree_id)) + export_tree_data = [ExportMessageNode.prep_message_export(m) for m in messages] + + message_parents = defaultdict(list) + for message in export_tree_data: + message_parents[message.parent_id].append(message) + + def build_tree(tree: dict, parent: Optional[str], messages: list[Message]): + children = message_parents[parent] + tree.replies = children + + for idx, child in enumerate(tree.replies): + build_tree(tree.replies[idx], child.message_id, messages) + + build_tree(export_tree, None, export_tree_data) + + return export_tree + + +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") + else: + out_buff = open(file, "wt", encoding="UTF-8") + + with out_buff as f: + for tree in trees: + file_data = jsonable_encoder(tree) + json.dump(file_data, f) + f.write("\n")