add new branch for multiple models

app/app.py

@@ -10,17 +10,29 @@ import psutil
 
 # st.set_page_config(page_title="Indonesian Story Generator")
 
-# vector_length = 128
-model_name = "cahya/gpt2-small-indonesian-story"
+MODELS = {
+    "GPT-2 Small finetuned on Indonesian stories": {
+        "name": "cahya/gpt2-small-indonesian-story"
+    },
+    "GPT-2 Medium finetuned on Indonesian stories": {
+        "name": "cahya/gpt2-medium-indonesian-story"
+    },
+}
+
+model = st.selectbox('Model',([
+    'GPT-2 Small finetuned on Indonesian stories',
+    'GPT-2 Medium finetuned on Indonesian stories']))
+
+if model in ["GPT-2 Small finetuned on Indonesian stories", "GPT-2 Medium finetuned on Indonesian stories"]:
+    prompt_group_name = "Indonesian Stories"
 
 
 @st.cache(suppress_st_warning=True, allow_output_mutation=True)
-def get_generator():
+def get_generator(model_name: str):
     st.write(f"Loading the GPT2 model {model_name}, please wait...")
     text_generator = pipeline('text-generation', model=model_name)
     return text_generator
 
-
 @st.cache(suppress_st_warning=True, hash_funcs={tokenizers.Tokenizer: id})
 def process(text: str, max_length: int = 100, do_sample: bool = True, top_k: int = 50, top_p: float = 0.95,
             temperature: float = 1.0, max_time: float = 10.0, seed=42):
@@ -39,9 +51,13 @@ st.markdown(
     This application is a demo for Indonesian Story Generator using GPT2.
     """
 )
+
+
 session_state = SessionState.get(prompt=None, prompt_box=None, text=None)
-ALL_PROMPTS = list(PROMPT_LIST.keys())+["Custom"]
+
+ALL_PROMPTS = list(PROMPT_LIST[prompt_group_name].keys())+["Custom"]
 prompt = st.selectbox('Prompt', ALL_PROMPTS, index=len(ALL_PROMPTS)-1)
+
 # Update prompt
 if session_state.prompt is None:
     session_state.prompt = prompt

app/prompts.py

@@ -1,18 +1,20 @@
 PROMPT_LIST = {
-    "Adult Romance": [
-        "Ini adalah kisah tentang seorang laki-laki yang berusaha memperjuangkan cintanya",
-        "Alunan musik terdengar memenuhi ruangan kantor, cowok itu duduk di balik meja kerjanya sambil memejamkan mata. Berusaha meresapi nada per nada",
-        "Aku mencari dan terus mencari\nDimana bahagia akan kutemui\nKumencari terus mencari\nHingga ku tak mengerti arti hari-hari",
-        "Gadis itu mengharuskan dirinya tegar, dan kuat dalam menghadapi masalah. Menahan air matanya jatuh setiap kali ingin menangis"
-    ],
-    "Horror": [
-        "Ditengah-tengah perbincangan mereka berdua, datanglah sesosok mahluk tinggi hitam dan besar",
-        "Sesosok hantu perempuan seperti kuntilanak yang melayang keluar dan bergerak perlahan dari pintu kamar kecil tadi yang tertutup.",
-        "Sejak pertemuannya dengan leak, yang ternyata tinggal satu atap dengannya, hidupnya terus dihantui oleh berbagai sosok seram."
-    ],
-    "Poetry": [
-        "Aku ingin menulis sajak\nyang melesat dalam kejap\nmenembus hati yang pejam\nmemaksa mimpimu terjaga\ndari semu",
-        "Malam ini langitku lengang\ntiada hujan yang membasuh rindu\npun awan yang biasanya temani seruput kopimu",
-        "Di sisimu waktu menjelma\nsetangkai kembang api\ngelora membakar tanpa jeda\nmemercik pijar binar kita."
-    ]
+    "Indonesian Stories": {
+        "Adult Romance": [
+            "Ini adalah kisah tentang seorang laki-laki yang berusaha memperjuangkan cintanya",
+            "Alunan musik terdengar memenuhi ruangan kantor, cowok itu duduk di balik meja kerjanya sambil memejamkan mata. Berusaha meresapi nada per nada",
+            "Aku mencari dan terus mencari\nDimana bahagia akan kutemui\nKumencari terus mencari\nHingga ku tak mengerti arti hari-hari",
+            "Gadis itu mengharuskan dirinya tegar, dan kuat dalam menghadapi masalah. Menahan air matanya jatuh setiap kali ingin menangis"
+        ],
+        "Horror": [
+            "Ditengah-tengah perbincangan mereka berdua, datanglah sesosok mahluk tinggi hitam dan besar",
+            "Sesosok hantu perempuan seperti kuntilanak yang melayang keluar dan bergerak perlahan dari pintu kamar kecil tadi yang tertutup.",
+            "Sejak pertemuannya dengan leak, yang ternyata tinggal satu atap dengannya, hidupnya terus dihantui oleh berbagai sosok seram."
+        ],
+        "Poetry": [
+            "Aku ingin menulis sajak\nyang melesat dalam kejap\nmenembus hati yang pejam\nmemaksa mimpimu terjaga\ndari semu",
+            "Malam ini langitku lengang\ntiada hujan yang membasuh rindu\npun awan yang biasanya temani seruput kopimu",
+            "Di sisimu waktu menjelma\nsetangkai kembang api\ngelora membakar tanpa jeda\nmemercik pijar binar kita."
+        ]
+    }
 }