Adding text-gen and fill-mask tasks

multipage.py

@@ -0,0 +1,40 @@
+"""
+This file is the framework for generating multiple Streamlit applications 
+through an object oriented framework. 
+"""
+
+# Import necessary libraries 
+import streamlit as st
+
+# Define the multipage class to manage the multiple apps in our program 
+class MultiPage: 
+    """Framework for combining multiple streamlit applications."""
+
+    def __init__(self) -> None:
+        """Constructor class to generate a list which will store all our applications as an instance variable."""
+        self.pages = []
+    
+    def add_page(self, title, func) -> None: 
+        """Class Method to Add pages to the project
+        Args:
+            title ([str]): The title of page which we are adding to the list of apps 
+            
+            func: Python function to render this page in Streamlit
+        """
+
+        self.pages.append({
+          
+                "title": title, 
+                "function": func
+            })
+
+    def run(self):
+        # Drodown to select the page to run  
+        page = st.sidebar.selectbox(
+            'App Navigation', 
+            self.pages, 
+            format_func=lambda page: page['title']
+        )
+
+        # run the app function 
+        page['function']()

pages/fill_mask.py

@@ -0,0 +1,64 @@
+import streamlit as st
+import time
+import pandas as pd
+import altair as alt
+from multipage import MultiPage
+from transformers import pipeline
+
+def app():
+    st.markdown('## Mask Fill task')
+    st.write('Write a sentence with a [MASK] gap to fill')
+    st.markdown('## ')
+    
+
+    @st.cache(allow_output_mutation=True, suppress_st_warning =True, show_spinner=False)
+    def get_model():
+        return pipeline('fill-mask', model='bert-base-uncased',skip_special_tokens=True)
+
+    def create_graph(answer):
+        x_bar = [i['token_str'] for i in answer]
+        y_bar = [i['score'] for i in answer]
+        chart_data = pd.DataFrame(y_bar, index=x_bar)
+        data = pd.melt(chart_data.reset_index(), id_vars=["index"])
+        # Horizontal stacked bar chart
+        chart = (
+            alt.Chart(data)
+            .mark_bar(color='#d7abf5')
+            .encode(
+                x=alt.X("index", type="nominal", title='',sort=alt.EncodingSortField(field="index", op="count", order='ascending')),
+                y=alt.Y("value", type="quantitative", title="Score", sort='-x'),
+            )
+        )
+        st.altair_chart(chart, use_container_width=True)
+
+        
+    col1, col2 = st.beta_columns([2,1])
+
+
+    with col1:
+        prompt= st.text_area('Your prompt here',
+            '''Who is Elon [MASK]?''') 
+            
+    with col2:
+        select_model = st.radio(
+            "Select the model to use:",
+            ('Bert cased', 'Bert Un-cased'), index = 1)
+
+        if select_model == 'Bert cased':
+            model = 'bert-base-cased'
+        elif select_model == 'Bert Un-cased':
+            model = 'bert-base-uncased'
+
+        with st.spinner('Loading Model... (This may take a while)'):
+            # unmasker = get_model()    
+            st.success('Model loaded correctly!')
+        
+            gen = st.info('Generating text...')
+            # answer = unmasker(prompt)     
+            answer = [{'sequence': "[CLS] hello i'm a fashion model. [SEP]",  'score': 0.1073106899857521,  'token': 4827,  'token_str': 'fashion'}, {'sequence': "[CLS] hello i'm a role model. [SEP]",  'score': 0.08774490654468536,  'token': 2535,  'token_str': 'role'}, {'sequence': "[CLS] hello i'm a new model. [SEP]",  'score': 0.05338378623127937,  'token': 2047,  'token_str': 'new'}, {'sequence': "[CLS] hello i'm a super model. [SEP]",  'score': 0.04667217284440994,  'token': 3565,  'token_str': 'super'}, {'sequence': "[CLS] hello i'm a fine model. [SEP]",  'score': 0.027095865458250046, 'token': 2986,  'token_str': 'fine'}]                 
+            gen.empty()    
+
+    with col1:   
+        create_graph(answer)       
+
+

pages/home_page.py

@@ -0,0 +1,9 @@
+import streamlit as st
+from multipage import MultiPage
+
+def app():
+    st.title('All-in-One')
+    st.markdown('#### Select a specific task from the sidebar...')
+
+
+

pages/text_gen.py

@@ -0,0 +1,53 @@
+import streamlit as st
+import time
+from multipage import MultiPage
+from transformers import pipeline
+import torch
+
+
+def app():
+    st.markdown('## Text Generation task')
+    st.write('Write something and AI will continue the sentence ')
+    st.markdown('## ')
+
+    @st.cache(allow_output_mutation=True, suppress_st_warning =True, show_spinner=False)
+    def get_model():
+        return pipeline('text-generation', model=model, do_sample=True, skip_special_tokens=True)
+        
+    col1, col2 = st.beta_columns([2,1])
+
+        
+    with col1:
+        prompt= st.text_area('Your prompt here',
+            '''Who is Elon Musk?''') 
+            
+    with col2:
+        select_model = st.radio(
+            "Select the model to use:",
+            ('OPT-125m', 'OPT-350m'), index = 1)
+        if select_model == 'OPT-350m':
+            model = 'facebook/opt-350m'
+        elif select_model == 'OPT-125m':
+            model = 'facebook/opt-125m'
+
+        with st.spinner('Loading Model... (This may take a while)'):
+            generator = get_model()    
+            st.success('Model loaded correctly!')
+
+    with col1:    
+        gen = st.info('Generating text...')
+        answer = generator(prompt,
+                                max_length=max_length, no_repeat_ngram_size=no_ngram_repeat,
+                                 early_stopping=early_stopping, num_beams=num_beams,
+                                   skip_special_tokens=True)                      
+        gen.empty()                      
+                        
+    lst = answer[0]['generated_text']
+    
+    t = st.empty()
+    for i in range(len(lst)):
+        t.markdown("#### %s..." % lst[0:i])
+        time.sleep(0.04)
+
+
+

requirements.txt

@@ -0,0 +1,5 @@
+streamlit
+transformers
+altair
+pandas
+torch