Add type annotation to Model.get()

2022-11-27 13:36:49 -08:00 · 2022-11-27 13:36:49 -08:00 · 8d42a92848
commit 8d42a92848
parent e4eb322aa7
1 changed files with 1 additions and 1 deletions
--- a/gptc/model.py
+++ b/gptc/model.py
@ -62,7 +62,7 @@ class Model:
        }
        return probs

-    def get(self, token):
+    def get(self, token: str) -> Dict[str, float]:
        try:
            weights = self.weights[
                gptc.tokenizer.hash_single(gptc.tokenizer.normalize(token))