fix: lru cache HF get model params (#105)

HazyResearch · Jul 2, 2023 · fd7fbc9 · fd7fbc9
1 parent b745617
commit fd7fbc9
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 0 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -1,5 +1,9 @@
 0.1.9 - Unreleased
 ---------------------
+Fixed
+^^^^^
+* Added trust code params HF models
+* Added LRU cache to HF model param calls to avoid extra calls
 
 0.1.8 - 2023-05-22
 ---------------------

diff --git a/manifest/clients/diffuser.py b/manifest/clients/diffuser.py
@@ -1,5 +1,6 @@
 """Diffuser client."""
 import logging
+from functools import lru_cache
 from typing import Any, Dict, Optional
 
 import numpy as np
@@ -79,6 +80,7 @@ def supports_streaming_inference(self) -> bool:
  """
  return False
 
+ @lru_cache(maxsize=1)
  def get_model_params(self) -> Dict:
  """
  Get model params.

diff --git a/manifest/clients/huggingface.py b/manifest/clients/huggingface.py
@@ -1,5 +1,6 @@
 """Hugging Face client."""
 import logging
+from functools import lru_cache
 from typing import Any, Dict, Optional
 
 import requests
@@ -73,6 +74,7 @@ def supports_streaming_inference(self) -> bool:
  """
  return False
 
+ @lru_cache(maxsize=1)
  def get_model_params(self) -> Dict:
  """
  Get model params.

diff --git a/manifest/clients/huggingface_embedding.py b/manifest/clients/huggingface_embedding.py
@@ -1,5 +1,6 @@
 """Hugging Face client."""
 import logging
+from functools import lru_cache
 from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
@@ -65,6 +66,7 @@ def supports_streaming_inference(self) -> bool:
  """
  return False
 
+ @lru_cache(maxsize=1)
  def get_model_params(self) -> Dict:
  """
  Get model params.