Pippoz commited on
Commit
df9431b
1 Parent(s): 429d718

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -7
app.py CHANGED
@@ -7,14 +7,14 @@ st.markdown('## Text-generation OPT from Meta ')
7
 
8
  @st.cache(allow_output_mutation=True)
9
  def get_model():
10
- return pipeline('text-generation', model=model)
11
 
12
  col1, col2 = st.beta_columns([2,1])
13
 
14
  with col2:
15
  select_model = st.radio(
16
  "Select the model to use:",
17
- ('OPT-125m', 'OPT-350m', 'OPT-1.3b'))
18
 
19
  if select_model == 'OPT-1.3b':
20
  model = 'facebook/opt-1.3b'
@@ -23,11 +23,9 @@ with col2:
23
  elif select_model == 'OPT-125m':
24
  model = 'facebook/opt-125m'
25
 
26
- if select_model:
27
- with st.spinner('Loading Model... (This may take a while)'):
28
- generator = get_model()
29
- #time.sleep(2)
30
- st.success('Model loaded correctly!')
31
 
32
 
33
  with col1:
 
7
 
8
  @st.cache(allow_output_mutation=True)
9
  def get_model():
10
+ return pipeline('text-generation', model=model, skip_special_tokens=True)
11
 
12
  col1, col2 = st.beta_columns([2,1])
13
 
14
  with col2:
15
  select_model = st.radio(
16
  "Select the model to use:",
17
+ ('OPT-125m', 'OPT-350m', 'OPT-1.3b'), index = 1)
18
 
19
  if select_model == 'OPT-1.3b':
20
  model = 'facebook/opt-1.3b'
 
23
  elif select_model == 'OPT-125m':
24
  model = 'facebook/opt-125m'
25
 
26
+ with st.warning('Loading Model... (This may take a while)'):
27
+ generator = get_model()
28
+ st.success('Model loaded correctly!')
 
 
29
 
30
 
31
  with col1: