use sets for shard specific patterns

e0fda94d2094903f03c32c908fd37b4274c5fab4 · 2024-09-05 15:58:41 +0100 · Alex Cheema

Files touched

M exo/download/hf/hf_helpers.py

Diff

commit e0fda94d2094903f03c32c908fd37b4274c5fab4
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Thu Sep 5 15:58:41 2024 +0100

    use sets for shard specific patterns
---
 exo/download/hf/hf_helpers.py | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/exo/download/hf/hf_helpers.py b/exo/download/hf/hf_helpers.py
index d4251388..d8742222 100644
--- a/exo/download/hf/hf_helpers.py
+++ b/exo/download/hf/hf_helpers.py
@@ -391,25 +391,19 @@ def extract_layer_num(tensor_name: str) -> Optional[int]:
 
 
 def get_allow_patterns(weight_map: Dict[str, str], shard: Shard) -> List[str]:
-  default_patterns = [
-    "*.json",
-    "*.py",
-    "tokenizer.model",
-    "*.tiktoken",
-    "*.txt",
-  ]
-  shard_specific_patterns = []
+  default_patterns = set(["*.json","*.py","tokenizer.model","*.tiktoken","*.txt"])
+  shard_specific_patterns = set()
   if weight_map:
     for tensor_name, filename in weight_map.items():
       layer_num = extract_layer_num(tensor_name)
       if layer_num is not None and shard.start_layer <= layer_num <= shard.end_layer:
-        shard_specific_patterns.append(filename)
+        shard_specific_patterns.add(filename)
     sorted_file_names = sorted(weight_map.values())
     if shard.is_first_layer():
-      shard_specific_patterns.append(sorted_file_names[0])
+      shard_specific_patterns.add(sorted_file_names[0])
     elif shard.is_last_layer():
-      shard_specific_patterns.append(sorted_file_names[-1])
+      shard_specific_patterns.add(sorted_file_names[-1])
   else:
     shard_specific_patterns = ["*.safetensors"]
   if DEBUG >= 2: print(f"get_allow_patterns {weight_map=} {shard=} {shard_specific_patterns=}")
-  return list(set(default_patterns + shard_specific_patterns))  # Remove duplicates
+  return list(default_patterns | shard_specific_patterns)  # Remove duplicates

← 8f65e1e6 fix weight_map resolution. previously we were always default · back to Exo · remove comment ea3322de →