fix

predibase · Dec 19, 2023 · 0054956 · 0054956
1 parent a1d01c6
commit 0054956
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/server/lorax_server/models/causal_lm.py b/server/lorax_server/models/causal_lm.py
@@ -533,7 +533,7 @@ def decode(self, generated_ids: List[int]) -> str:
         )
 
     def forward(
-        self, input_ids, attention_mask, position_ids, past_key_values: Optional = None
+        self, input_ids, attention_mask, position_ids, past_key_values: Optional[torch.Tensor] = None
     ) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]:
         # Model Forward
         kwargs = {

diff --git a/server/lorax_server/utils/sources/hub.py b/server/lorax_server/utils/sources/hub.py
@@ -11,6 +11,7 @@
 from huggingface_hub.utils import (
     LocalEntryNotFoundError,
     EntryNotFoundError,  # Import here to ease try/except in other part of the lib
+    RevisionNotFoundError,
 )
 
 from .source import BaseModelSource, try_to_load_from_cache