From 8d42a92848038d5f7caf72b6c3f8d0d39b38bb11 Mon Sep 17 00:00:00 2001
From: Samuel Sloniker <sam@kj7rrv.com>
Date: Sun, 27 Nov 2022 13:36:49 -0800
Subject: [PATCH] Add type annotation to Model.get()

---
 gptc/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gptc/model.py b/gptc/model.py
index 634d900..53e7bbc 100644
--- 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))