From 2d37a7341ae83c6a69ab66372ed7fdd5570c3052 Mon Sep 17 00:00:00 2001 From: pixelpaws Date: Fri, 3 Oct 2025 22:45:25 +0800 Subject: [PATCH] fix(routes): await trained words extraction --- py/routes/handlers/misc_handlers.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/py/routes/handlers/misc_handlers.py b/py/routes/handlers/misc_handlers.py index de553fa6..79b581e9 100644 --- a/py/routes/handlers/misc_handlers.py +++ b/py/routes/handlers/misc_handlers.py @@ -320,9 +320,14 @@ class TrainedWordsHandler: if not file_path.endswith(".safetensors"): return web.json_response({"success": False, "error": "File must be a safetensors file"}, status=400) - trained_words = extract_trained_words(file_path) - sorted_words = sorted(trained_words, key=lambda w: w.get("count", 0), reverse=True) - return web.json_response({"success": True, "trained_words": sorted_words}) + trained_words, class_tokens = await extract_trained_words(file_path) + return web.json_response( + { + "success": True, + "trained_words": trained_words, + "class_tokens": class_tokens, + } + ) except Exception as exc: # pragma: no cover - defensive logging logger.error("Failed to get trained words: %s", exc, exc_info=True) return web.json_response({"success": False, "error": str(exc)}, status=500)