Update pages/fill_mask.py

pages/fill_mask.py

@@ -12,8 +12,8 @@ def app():
     
 
     @st.cache(allow_output_mutation=True, suppress_st_warning =True, show_spinner=False)
-    def get_model():
-        return pipeline('fill-mask', model='bert-base-uncased')
+    def get_model(model):
+        return pipeline('fill-mask', model=model)
 
     def create_graph(answer):
         x_bar = [i['token_str'] for i in answer]
@@ -50,7 +50,7 @@ def app():
             model = 'bert-base-uncased'
 
         with st.spinner('Loading Model... (This may take a while)'):
-            unmasker = get_model()    
+            unmasker = get_model(model)    
             st.success('Model loaded correctly!')
         
             gen = st.info('Generating text...')