idiomify / main_deploy.py
eubinecto's picture
[#2] minor changes to deploy settings
70bfe2b
raw
history blame
1.5 kB
"""
we deploy the pipeline via streamlit.
"""
from typing import Tuple, List
import streamlit as st
from transformers import BartTokenizer
from idiomify.fetchers import fetch_config, fetch_idiomifier, fetch_idioms
from idiomify.pipeline import Pipeline
from idiomify.models import Idiomifier
@st.cache(allow_output_mutation=True)
def fetch_resources() -> Tuple[dict, Idiomifier, BartTokenizer, List[str]]:
config = fetch_config()['idiomifier']
model = fetch_idiomifier(config['ver'])
idioms = fetch_idioms(config['idioms_ver'])
tokenizer = BartTokenizer.from_pretrained(config['bart'])
return config, model, tokenizer, idioms
def main():
# fetch a pre-trained model
config, model, tokenizer, idioms = fetch_resources()
pipeline = Pipeline(model, tokenizer)
st.title("Idiomify Demo")
st.markdown(f"Author: `Eu-Bin KIM`")
st.markdown(f"Version: `{config['ver']}`")
text = st.text_area("Type sentences here",
value="Just remember there will always be a hope even when things look black")
with st.sidebar:
st.subheader("Supported idioms")
st.write(" / ".join(idioms))
if st.button(label="Idiomify"):
with st.spinner("Please wait..."):
sents = [sent for sent in text.split(".") if sent]
sents = pipeline(sents, max_length=200)
# highlight the rule & honorifics that were applied
st.write(". ".join(sents))
if __name__ == '__main__':
main()