Switch to AutoModels

app.py

@@ -11,7 +11,7 @@ import numpy as np
 import tensorflow as tf
 import streamlit as st
 from gazpacho import Soup, get
-from transformers import BertTokenizer, TFBertForMaskedLM
+from transformers import AutoTokenizer, TFAutoModelForMaskedLM
 
 
 DEFAULT_QUERY = "Machines will take over the world soon"
@@ -85,8 +85,8 @@ def start_rhyming(query, rhyme_words_options):
 @st.cache(allow_output_mutation=True)
 def load_model(model_path):
     return (
-        TFBertForMaskedLM.from_pretrained(model_path),
-        BertTokenizer.from_pretrained(model_path),
+        TFAutoModelForMaskedLM.from_pretrained(model_path),
+        AutoTokenizer.from_pretrained(model_path),
     )
 
 
@@ -119,8 +119,8 @@ class TokenWeighter:
 class RhymeGenerator:
     def __init__(
         self,
-        model: TFBertForMaskedLM,
-        tokenizer: BertTokenizer,
+        model: TFAutoModelForMaskedLM,
+        tokenizer: AutoTokenizer,
         token_weighter: TokenWeighter = None,
     ):
         """Generate rhymes.