use set for shard specific patterns

11dd952d26f17c241fb7da79a1a28342e2f951d3 · 2024-09-05 16:00:59 +0100 · Alex Cheema

Files touched

M exo/download/hf/hf_helpers.py

Diff

commit 11dd952d26f17c241fb7da79a1a28342e2f951d3
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Thu Sep 5 16:00:59 2024 +0100

    use set 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..3197605d 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)

← ea3322de remove comment · back to Exo · add llama-3.1-70b-bf16 model option 2948a834 →