Update app.py

app.py

@@ -3,20 +3,45 @@ import time
 from transformers import pipeline
 import torch
 
-st.markdown('## OPT-1.3 Billion parameter (Meta)')
+st.markdown('## Text-generation OPT from Meta ')
 
-with st.spinner('Loading Model... (This may take a while)'):
-    generator = pipeline('text-generation', model="facebook/opt-1.3b", skip_special_tokens=True)    
-    st.success('Model loaded correctly!')
+@st.cache(allow_output_mutation=True)
+def get_model():
+    return pipeline('text-generation', model=model)
 
-prompt= st.text_area('Your prompt here',
- '''AI will help humanity?''') 
+col1, col2 = st.beta_columns([2,1])
+
+with col2:
+    select_model = st.radio(
+        "Select the model to use:",
+        ('OPT-125m', 'OPT-350m', 'OPT-1.3b'))
+
+    if select_model == 'OPT-1.3b':
+        model = 'facebook/opt-1.3b'
+    elif select_model == 'OPT-350m':
+        model = 'facebook/opt-350m'
+    elif select_model == 'OPT-125m':
+        model = 'facebook/opt-125m'
+
+    if select_model:
+        with st.spinner('Loading Model... (This may take a while)'):
+            generator = get_model()    
+            #time.sleep(2)
+            st.success('Model loaded correctly!')
+    
+
+with col1:
+    prompt= st.text_area('Your prompt here',
+        '''AI will help humanity?''') 
+
+    # answer = ['ciao come stai stutto bene']
+    # lst = ''.join(answer)
  
 answer = generator(prompt, max_length=60, no_repeat_ngram_size=2, early_stopping=True, num_beams=5)
- 
 lst = answer[0]['generated_text']
 
-t = st.empty()
-for i in range(len(lst)):
-    t.markdown(" %s..." % lst[0:i])
-    time.sleep(0.04)
+
+    t = st.empty()
+    for i in range(len(lst)):
+        t.markdown(" %s..." % lst[0:i])
+        time.sleep(0.04)