From d5b7e6e38fa4f4fa4e1c7c34367d9f8069393367 Mon Sep 17 00:00:00 2001 From: Raphael Glon Date: Fri, 15 Dec 2023 12:48:05 +0100 Subject: [PATCH] Reuse the same function to list local weights everywhere Signed-off-by: Raphael Glon --- server/text_generation_server/utils/hub.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/server/text_generation_server/utils/hub.py b/server/text_generation_server/utils/hub.py index b93eeeac..019d4855 100644 --- a/server/text_generation_server/utils/hub.py +++ b/server/text_generation_server/utils/hub.py @@ -125,13 +125,14 @@ def weight_files( ) -> List[Path]: """Get the local files""" # Local model - if Path(model_id).exists() and Path(model_id).is_dir(): - local_files = list(Path(model_id).glob(f"*{extension}")) + d = Path(model_id) + if d.exists() and d.is_dir(): + local_files = _weight_files_from_dir(d, extension) if not local_files: raise FileNotFoundError( f"No local weights found in {model_id} with extension {extension}" ) - return local_files + return [Path(f) for f in local_files] try: filenames = weight_hub_files(model_id, revision, extension)