Create token_weighter.py

rhyme_with_ai/token_weighter.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+
+class TokenWeighter:
+    def __init__(self, tokenizer):
+        self.tokenizer_ = tokenizer
+        self.proba = self.get_token_proba()
+
+    def get_token_proba(self):
+        valid_token_mask = self._filter_short_partial(self.tokenizer_.vocab)
+        return valid_token_mask
+
+    def _filter_short_partial(self, vocab):
+        valid_token_ids = [v for k, v in vocab.items() if len(k) > 1 and "#" not in k]
+        is_valid = np.zeros(len(vocab.keys()))
+        is_valid[valid_token_ids] = 1
+        return is_valid