Adding finetuned model as a demo

app.py

@@ -1,28 +1,59 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+# from huggingface_hub import snapshot_download
 
-st.title('Sinhala Text generation with GPT2')
-st.markdown('A simple demo using Sinhala-gpt2 model trained during hf-flax week')
+page = st.sidebar.selectbox("Model ", ["Pretrained GPT2", "Finetuned on News data"])
 
-seed = st.text_input('Starting text', 'ආයුබෝවන්')
-seq_num = st.number_input('Number of sentences to generate ', 1, 20, 5)
-max_len = st.number_input('Length of the sentence ', 5, 300, 100)
+def load_model(model_name):
+    with st.spinner('Waiting for the model to load.....'):
+        # snapshot_download('flax-community/Sinhala-gpt2')
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+    st.success('Model loaded!!')
+    return model, tokenizer
 
-go = st.button('Generate')
-with st.spinner('Waiting for the model to load.....'):
-    model = AutoModelForCausalLM.from_pretrained('flax-community/Sinhala-gpt2')
-    tokenizer = AutoTokenizer.from_pretrained('flax-community/Sinhala-gpt2')
-st.success('Model loaded!!')
+seed = st.sidebar.text_input('Starting text', 'ආයුබෝවන්')
+seq_num = st.sidebar.number_input('Number of sentences to generate ', 1, 20, 5)
+max_len = st.sidebar.number_input('Length of the sentence ', 5, 300, 100)
 
+if page == "Finetuned on News data":
+    
+    st.title('Sinhala Text generation with Finetuned GPT2')
+    st.markdown('This model has been finetuned Sinhala-gpt2 model with 6000 news articles(~12MB)')
+    
+    # seed = st.text_input('Starting text', 'ආයුබෝවන්')
+    # seq_num = st.number_input('Number of sentences to generate ', 1, 20, 5)
+    # max_len = st.number_input('Length of the sentence ', 5, 300, 100)
 
-if go:
-    try:
-        with st.spinner('Generating...'):
-            generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
-            seqs = generator(seed, max_length=max_len, num_return_sequences=seq_num)
-        st.write(seqs)
-    except Exception as e:
-        st.exception(f'Exception: {e}')
+    gen_news = st.button('Generate')
+    model, tokenizer = load_model('keshan/sinhala-gpt2-newswire')
+
+
+    if gen_news:
+        try:
+            with st.spinner('Generating...'):
+                generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
+                seqs = generator(seed, max_length=max_len, num_return_sequences=seq_num)
+            st.write(seqs)
+        except Exception as e:
+            st.exception(f'Exception: {e}')
+else:
+    st.title('Sinhala Text generation with GPT2')
+    st.markdown('A simple demo using Sinhala-gpt2 model trained during hf-flax week')
+
+    gen_gpt2 = st.button('Generate')
+    model, tokenizer = load_model('flax-community/Sinhala-gpt2')
+
+
+    if gen_gpt2:
+        try:
+            with st.spinner('Generating...'):
+                generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
+                seqs = generator(seed, max_length=max_len, num_return_sequences=seq_num)
+            st.write(seqs)
+        except Exception as e:
+            st.exception(f'Exception: {e}')
+            
 
 st.markdown('____________')
 st.markdown('by Keshan with Flax Community')