Repository: LAION-AI/Open-Assistant
File: inference/server/export.py (Python)

Code before:
        if not file:
            raise RuntimeError("File name must be specified when using compression.")
        out_buff = gzip.open(file, "wt", encoding="UTF-8")
    else:
        out_buff = smart_open(file)

    with out_buff as f:
        for chat in chats:
            if write_trees:
                export_chat = prepare_export_message_tree(chat, anonymizer=anonymizer)
                file_data = jsonable_encoder(export_chat, exclude_none=True)
                json.dump(file_data, f)
                f.write("\n")
            else:
                # Exclude messages without content (e.g. work still in progress or aborted)
                for message in filter(lambda m: m.content, chat.messages):
                    export_message = prepare_export_message_node(chat, message, anonymizer=anonymizer)
                    file_data = jsonable_encoder(export_message, exclude_none=True)
                    json.dump(file_data, f)
                    f.write("\n")


async def fetch_eligible_chats(
    session_generator, filters: list[sqlalchemy.sql.elements.BinaryExpression]
) -> list[DbChat]:
    """Fetch chats which are not opted out of data collection and match the given filters."""
    session: AsyncSession
    filters.append(DbChat.allow_data_use)
    async with session_generator() as session:
        query = (
            sqlmodel.select(DbChat)
            .filter(*filters)
            .options(
                sqlalchemy.orm.joinedload("*"),
            )
        )
        chats: list[DbChat] = (await session.exec(query)).unique().all()
        return chats


def export_chats(
    session_generator,
    export_path: Path,
    filters: list[sqlalchemy.sql.elements.BinaryExpression],
    use_compression: bool = True,
    write_trees: bool = True,
    anonymizer_seed: str | None = None,
) -> None:
    eligible_chats: list[DbChat] = asyncio.run(fetch_eligible_chats(session_generator, filters))
    anonymizer = Anonymizer(anonymizer_seed) if anonymizer_seed else None


Review comment (suggestion):
This isn't really to do with your PR, but could you also add `not_(DbChat.deleted)` here? I forgot to include it when originally developing this script

Code after:
        if not file:
            raise RuntimeError("File name must be specified when using compression.")
        out_buff = gzip.open(file, "wt", encoding="UTF-8")
    else:
        out_buff = smart_open(file)

    with out_buff as f:
        for chat in chats:
            if write_trees:
                export_chat = prepare_export_message_tree(chat, anonymizer=anonymizer)
                file_data = jsonable_encoder(export_chat, exclude_none=True)
                json.dump(file_data, f)
                f.write("\n")
            else:
                # Exclude messages without content (e.g. work still in progress or aborted)
                for message in filter(lambda m: m.content, chat.messages):
                    export_message = prepare_export_message_node(chat, message, anonymizer=anonymizer)
                    file_data = jsonable_encoder(export_message, exclude_none=True)
                    json.dump(file_data, f)
                    f.write("\n")


async def fetch_eligible_chats(session_generator, filters: list[Any]) -> list[DbChat]:
    """Fetch chats which are not opted out of data collection and match the given filters."""
    session: AsyncSession
    filters.append(DbChat.allow_data_use)
    async with session_generator() as session:
        query = (
            sqlmodel.select(DbChat)
            .filter(*filters)
            .options(
                sqlalchemy.orm.joinedload("*"),
            )
        )
        chats: list[DbChat] = (await session.exec(query)).unique().all()
        return chats


def export_chats(
    session_generator,
    export_path: Path,
    filters: list[Any],
    use_compression: bool = True,
    write_trees: bool = True,
    anonymizer_seed: str | None = None,
) -> None:
    eligible_chats: list[DbChat] = asyncio.run(fetch_eligible_chats(session_generator, filters))
    anonymizer = Anonymizer(anonymizer_seed) if anonymizer_seed else None

    write_messages_to_file(
        export_path,