Simon Salmon commited on
Commit
e129d30
·
1 Parent(s): fde12d7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -15,8 +15,9 @@ st.title('KoGPT2 Demo')
15
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
16
 
17
  from transformers import AutoTokenizer, AutoModelForMaskedLM
 
18
  tokenizer = AutoTokenizer.from_pretrained("BigSalmon/FormalBerta")
19
- model = AutoModelForMaskedLM.from_pretrained("BigSalmon/FormalBerta")
20
 
21
  with st.form(key='my_form'):
22
  prompt = st.text_area(label='Enter sentence')
 
15
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
16
 
17
  from transformers import AutoTokenizer, AutoModelForMaskedLM
18
+ artist_name = st.text_input("Model", "BigSalmon/FormalBerta")
19
  tokenizer = AutoTokenizer.from_pretrained("BigSalmon/FormalBerta")
20
+ model = AutoModelForMaskedLM.from_pretrained(artist_name)
21
 
22
  with st.form(key='my_form'):
23
  prompt = st.text_area(label='Enter sentence')