🔧 refactor(infer.py): update type hints for FastText model references

LlmKira · Sep 29, 2024 · a52eb53 · a52eb53
1 parent d35775c
commit a52eb53
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/fast_langdetect/ft_detect/infer.py b/src/fast_langdetect/ft_detect/infer.py
@@ -32,10 +32,10 @@ class ModelCache:
     def __init__(self):
         self._models = {}
 
-    def get_model(self, model_type: ModelType) -> Optional[fasttext.FastText]:
+    def get_model(self, model_type: ModelType) -> Optional["fasttext.FastText._FastText"]:
         return self._models.get(model_type)
 
-    def set_model(self, model_type: ModelType, model: fasttext.FastText):
+    def set_model(self, model_type: ModelType, model: "fasttext.FastText._FastText"):
         self._models[model_type] = model