app.py

import streamlit as st
from huggingface_hub import snapshot_download

import os  # utility library

# libraries to load the model and serve inference
import tensorflow_text
import tensorflow as tf


def main():
    st.title("Interactive demo: T5 Multitasking Demo")
    st.sidebar.image("https://i.gzn.jp/img/2020/02/25/google-ai-t5/01.png")
    saved_model_path = load_model_cache()

    # Model is loaded in st.session_state to remain stateless across reloading
    if 'model' not in st.session_state:
        st.session_state.model = tf.saved_model.load(saved_model_path, ["serve"])

    dashboard(st.session_state.model)


@st.cache
def load_model_cache():
    """Function to retrieve the model from HuggingFace Hub and cache it using st.cache wrapper
    """
    CACHE_DIR = "hfhub_cache"  # where the library's fork would be stored once downloaded
    if not os.path.exists(CACHE_DIR):
        os.mkdir(CACHE_DIR)

    # download the files from huggingface repo and load the model with tensorflow
    snapshot_download(repo_id="stevekola/T5", cache_dir=CACHE_DIR)
    saved_model_path = os.path.join(CACHE_DIR, os.listdir(CACHE_DIR)[0])
    return saved_model_path


def dashboard(model):
    """Function to display the inputs and results
    params:
        model   stateless model to run inference from
    """
    task_type = st.sidebar.radio("Task Type",
                                 [
                                     "Translate English to French",
                                     "Translate English to German",
                                     "Translate English to Romanian",
                                     "Grammatical Correctness of Sentence",
                                     "Text Summarization",
                                     "Document Similarity Score"
                                 ])

    default_sentence = "I am Steven and I live in Lagos, Nigeria."
    text_summarization_sentence = "I don't care about those doing the comparison, but comparing \
        the Ghanaian Jollof Rice to Nigerian Jollof Rice is an insult to Nigerians."
    doc_similarity_sentence1 = "I reside in the commercial capital city of Nigeria, which is Lagos."
    doc_similarity_sentence2 = "I live in Lagos."
    help_msg = "You could either type in the sentences to run inferences on or use the upload button to \
        upload text files containing those sentences. The input sentence box, by default, displays sample \
        texts or the texts in the files that you've uploaded. Feel free to erase them and type in new sentences."

    if task_type.startswith("Document Similarity"):  # document similarity requires two documents
        uploaded_file = upload_files(help_msg, text="Upload 2 documents for similarity check", accept_multiple_files=True)
        if uploaded_file:
            sentence1 = st.text_area("Enter first document/sentence", uploaded_file[0], help=help_msg)
            sentence2 = st.text_area("Enter second document/sentence", uploaded_file[1], help=help_msg)
        else:
            sentence1 = st.text_area("Enter first document/sentence", doc_similarity_sentence1)
            sentence2 = st.text_area("Enter second document/sentence", doc_similarity_sentence2)
        sentence = sentence1 + "---" + sentence2  # to be processed like other tasks' single sentences
    else:
        uploaded_file = upload_files(help_msg)
        if uploaded_file:
            sentence = st.text_area("Enter sentence", uploaded_file, help=help_msg)
        elif task_type.startswith("Text Summarization"):  # text summarization's default input should be longer
            sentence = st.text_area("Enter sentence", text_summarization_sentence, help=help_msg)
        else:
            sentence = st.text_area("Enter sentence", default_sentence, help=help_msg)

    st.write("**Output Text**")
    with st.spinner("Waiting for prediction..."):  # spinner while model is running inferences
        output_text = predict(task_type, sentence, model)
        st.write(output_text)
        try:  # to workaround the environment's Streamlit version
            st.download_button("Download output text", output_text)
        except AttributeError:
            st.text("File download not enabled for this Streamlit version \U0001F612")


def upload_files(help_msg, text="Upload a text file here", accept_multiple_files=False):
    """Function to upload text files and return as string text
    params:
        text                    Display label for the upload button
        accept_multiple_files   params for the file_uploader function to accept more than a file
    returns:
        a string or a list of strings (in case of multiple files being uploaded)
    """

    def upload():
        uploaded_files = st.file_uploader(label="Upload text files only", 
                                          type="txt", help=help_msg,
                                          accept_multiple_files=accept_multiple_files)
        if st.button("Process"):
            if not uploaded_files:
                st.write("**No file uploaded!**")
                return None
            st.write("**Upload successful!**")
            if type(uploaded_files) == list:
                return [f.read().decode("utf-8") for f in uploaded_files]
            return uploaded_files.read().decode("utf-8")

    try:  # to workaround the environment's Streamlit version
        with st.expander(text):
            return upload()
    except AttributeError:
        return upload()


def predict(task_type, sentence, model):
    """Function to parse the user inputs, run the parsed text through the
    model and return output in a readable format.
    params:
        task_type   sentence representing the type of task to run on T5 model
        sentence    sentence to get inference on
        model       model to get inferences from
    returns:
        text decoded into a human-readable format.
    """
    task_dict = {
        "Translate English to French": "Translate English to French",
        "Translate English to German": "Translate English to German",
        "Translate English to Romanian": "Translate English to Romanian",
        "Grammatical Correctness of Sentence": "cola sentence",
        "Text Summarization": "summarize",
        "Document Similarity Score": "stsb",
    }
    question = f"{task_dict[task_type]}: {sentence}"  # parsing the user inputs into a format recognized by T5
    # Document Similarity takes in two sentences so it has to be parsed in a separate manner
    if task_type.startswith("Document Similarity"):
        sentences = sentence.split('---')
        question = f"{task_dict[task_type]} sentence1: {sentences[0]} sentence2: {sentences[1]}"
    return predict_fn([question], model)[0].decode('utf-8')


def predict_fn(x, model):
    """Function to get inferences from model on live data points.
    params:
        x       input text to run get output on
        model   model to run inferences from
    returns:
        a numpy array representing the output
    """
    return model.signatures['serving_default'](tf.constant(x))['outputs'].numpy()


if __name__ == "__main__":
    main()