diff --git a/commons/reward_model/models.py b/commons/reward_model/models.py index 77b341ba..78e70d80 100644 --- 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: