Tokenizers: Updated huggingface_models.py to support Safetensors mode…

…ls as well as pytorch (#2880) * Updated huggingface_models.py to support Safetensors models as well as pytorch --------- Co-authored-by: Frank Liu <[email protected]>
deepjavalibrary · Nov 29, 2023 · c057898 · c057898
1 parent c62a0e9
commit c057898
Showing 1 changed file with 12 additions and 9 deletions.
diff --git a/extensions/tokenizers/src/main/python/huggingface_models.py b/extensions/tokenizers/src/main/python/huggingface_models.py
@@ -60,17 +60,20 @@ def list_models(self, args: Namespace) -> List[dict]:
 
         api = HfApi()
         if args.model_name:
-            models = api.list_models(filter="pytorch",
-                                     search=args.model_name,
-                                     sort="downloads",
-                                     direction=-1,
-                                     limit=args.limit)
+            all_models = api.list_models(search=args.model_name,
+                                         sort="downloads",
+                                         direction=-1,
+                                         limit=args.limit)
             import_all = True
         else:
-            models = api.list_models(filter=f"{args.category},pytorch",
-                                     sort="downloads",
-                                     direction=-1,
-                                     limit=args.limit)
+            all_models = api.list_models(filter=args.category,
+                                         sort="downloads",
+                                         direction=-1,
+                                         limit=args.limit)
+        models = [
+            model for model in all_models
+            if 'pytorch' in model.tags or 'safetensors' in model.tags
+        ]
         if not models:
             if args.model_name:
                 logging.warning(f"no model found: {args.model_name}.")