Update registry.py

explosion · Oct 24, 2023 · caecd85 · caecd85
1 parent 63c7183
commit caecd85
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/spacy_llm/models/bedrock/registry.py b/spacy_llm/models/bedrock/registry.py
@@ -1,15 +1,15 @@
-from typing import Any, Callable, Dict, Iterable
+from typing import Any, Callable, Dict, Iterable, List
 
 from confection import SimpleFrozenDict
 
 from ...registry import registry
 from .model import Bedrock, Models
 
-_DEFAULT_RETRIES = 5
-_DEFAULT_TEMPERATURE = 0.0
-_DEFAULT_MAX_TOKEN_COUNT = 512
-_DEFAULT_TOP_P = 1
-_DEFAULT_STOP_SEQUENCES = []
+_DEFAULT_RETRIES: int = 5
+_DEFAULT_TEMPERATURE: float = 0.0
+_DEFAULT_MAX_TOKEN_COUNT: int = 512
+_DEFAULT_TOP_P: int = 1
+_DEFAULT_STOP_SEQUENCES: List[str] = []
 
 @registry.llm_models("spacy.Bedrock.Titan.Express.v1")
 def titan_express(