Spaces:
Runtime error
Runtime error
File size: 1,343 Bytes
7d5081f 4946d76 7d5081f 8f074bc c9ee852 8f074bc df9431b 8aa5b8d 378eb4f 8f074bc df9431b 8f074bc 378eb4f 217c852 df9431b 8f074bc 378eb4f 429d718 378eb4f bca744c 378eb4f 429d718 8f074bc 429d718 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import streamlit as st
import time
from transformers import pipeline
import torch
st.markdown('## Text-generation OPT from Meta ')
@st.cache(allow_output_mutation=True)
def get_model():
return pipeline('text-generation', model=model, skip_special_tokens=True)
col1, col2 = st.columns([2,1])
with col2:
select_model = st.radio(
"Select the model to use:",
('OPT-125m', 'OPT-350m', 'OPT-1.3b'), index = 1)
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:
st.warning('Loading Model... (This may take a while)')
generator = get_model()
st.success('Model loaded correctly!')
with col1:
prompt= st.text_area('Your prompt here',
'''AI will help humanity? ''')
answer = generator(prompt,
max_length=60,
no_repeat_ngram_size=2, early_stopping=True, num_beams=5, skip_special_tokens=True)
lst = answer[0]['generated_text']
# answer = ['ciao come stai stutto bene']
# lst = ''.join(answer)
t = st.empty()
for i in range(len(lst)):
t.markdown(" %s..." % lst[0:i])
time.sleep(0.04) |