From 1430544f0f26c350586ca6440ea059c477ca5d3f Mon Sep 17 00:00:00 2001 From: Roman Zlobin Date: Fri, 18 Oct 2024 17:24:15 +0300 Subject: [PATCH] fix context size calculation --- chatsky/utils/db_benchmark/basic_config.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/chatsky/utils/db_benchmark/basic_config.py b/chatsky/utils/db_benchmark/basic_config.py index 3941f7f12..cc22a11ee 100644 --- a/chatsky/utils/db_benchmark/basic_config.py +++ b/chatsky/utils/db_benchmark/basic_config.py @@ -150,14 +150,21 @@ async def info(self): - "misc_size" -- size of a misc field of a context. - "message_size" -- size of a misc field of a message. """ + def remove_db_from_context(ctx: Context): + ctx._storage = None + ctx.requests._storage = None + ctx.responses._storage = None + ctx.labels._storage = None + + starting_context = await get_context(MemoryContextStorage(), self.from_dialog_len, self.message_dimensions, self.misc_dimensions) + final_contex = await get_context(MemoryContextStorage(), self.to_dialog_len, self.message_dimensions, self.misc_dimensions) + remove_db_from_context(starting_context) + remove_db_from_context(final_contex) return { "params": self.model_dump(), "sizes": { - "starting_context_size": naturalsize(asizeof.asizeof(await self.get_context(MemoryContextStorage())), gnu=True), - "final_context_size": naturalsize( - asizeof.asizeof(await get_context(MemoryContextStorage(), self.to_dialog_len, self.message_dimensions, self.misc_dimensions)), - gnu=True, - ), + "starting_context_size": naturalsize(asizeof.asizeof(starting_context.model_dump(mode="json")), gnu=True), + "final_context_size": naturalsize(asizeof.asizeof(final_contex.model_dump(mode="json")), gnu=True), "misc_size": naturalsize(asizeof.asizeof(get_dict(self.misc_dimensions)), gnu=True), "message_size": naturalsize(asizeof.asizeof(get_message(self.message_dimensions)), gnu=True), },