fix: increase lru cache size to accomodate different models

tensorplex-labs · Feb 25, 2024 · 59ad2ce · 59ad2ce
1 parent 5f8e9b5
commit 59ad2ce
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/commons/reward_model/models.py b/commons/reward_model/models.py
@@ -35,13 +35,16 @@ class ModelZoo(StrEnum):
 
 
 # adjust cache maxsize depending on however many models you have in your ModelZoo
-@lru_cache(maxsize=10)
+max_num_miners = 256
+
+
+@lru_cache(maxsize=max_num_miners)
 def get_cached_model(model_name: ModelName):
     model = AutoModelForSequenceClassification.from_pretrained(model_name).eval()
     return model
 
 
-@lru_cache(maxsize=10)
+@lru_cache(maxsize=max_num_miners)
 def get_cached_tokenizer(
     model_name: ModelName,
 ) -> PreTrainedTokenizerFast | PreTrainedTokenizer: