Fix old type hints

app.py

@@ -11,8 +11,6 @@ import numpy as np
 import tensorflow as tf
 import streamlit as st
 from gazpacho import Soup, get
-from transformers.modeling_tf_bert import TFBertForPreTraining
-from transformers.tokenization_bert import PreTrainedTokenizer
 from transformers import BertTokenizer, TFBertForMaskedLM
 
 
@@ -102,12 +100,27 @@ def display_output(status_text, query, current_sentences, previous_sentences):
         query + ",<br>" + "".join(print_sentences), unsafe_allow_html=True
     )
 
+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
+
 
 class RhymeGenerator:
     def __init__(
         self,
-        model: TFBertForPreTraining,
-        tokenizer: PreTrainedTokenizer,
+        model: TFBertForMaskedLM,
+        tokenizer: BertTokenizer,
         token_weighter: TokenWeighter = None,
     ):
         """Generate rhymes.
@@ -335,26 +348,6 @@ def mick_rijmwoordenboek(word: str, n_words: int):
     return random.sample(results, min(len(results), n_words))
 
 
-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
-
-
-
 def color_new_words(new: str, old: str, color: str = "#eefa66") -> str:
     """Color new words in strings with a span."""