hgrif commited on
Commit
f9afcb3
β€’
1 Parent(s): 3aa3b62

Switch to BertTokenizer

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -11,7 +11,7 @@ import numpy as np
11
  import tensorflow as tf
12
  import streamlit as st
13
  from gazpacho import Soup, get
14
- from transformers import AutoTokenizer, TFAutoModelForMaskedLM
15
 
16
 
17
  DEFAULT_QUERY = "Machines will take over the world soon"
@@ -87,7 +87,7 @@ def start_rhyming(query, rhyme_words_options):
87
  def load_model(model_path):
88
  return (
89
  TFAutoModelForMaskedLM.from_pretrained(model_path),
90
- AutoTokenizer.from_pretrained(model_path),
91
  )
92
 
93
 
@@ -121,7 +121,7 @@ class RhymeGenerator:
121
  def __init__(
122
  self,
123
  model: TFAutoModelForMaskedLM,
124
- tokenizer: AutoTokenizer,
125
  token_weighter: TokenWeighter = None,
126
  ):
127
  """Generate rhymes.
 
11
  import tensorflow as tf
12
  import streamlit as st
13
  from gazpacho import Soup, get
14
+ from transformers import BertTokenizer, TFAutoModelForMaskedLM
15
 
16
 
17
  DEFAULT_QUERY = "Machines will take over the world soon"
 
87
  def load_model(model_path):
88
  return (
89
  TFAutoModelForMaskedLM.from_pretrained(model_path),
90
+ BertTokenizer.from_pretrained(model_path),
91
  )
92
 
93
 
 
121
  def __init__(
122
  self,
123
  model: TFAutoModelForMaskedLM,
124
+ tokenizer: BertTokenizer,
125
  token_weighter: TokenWeighter = None,
126
  ):
127
  """Generate rhymes.