chose by checkpoint

app.py

@@ -30,18 +30,27 @@ def generation_markup_format(text):
 ds = load_dataset("bigscience/bloom-generations", use_auth_token=HF_API_TOKEN)
 ds = ds["train"]
 
-possible_langs = ds.unique("lang")
 
 col_1, col_2 = st.columns(2)
 with col_1:
+    possible_checkpoint = ds.unique("checkpoint")
     st.markdown("<h1 style='text-align: center'>Prompt</h1>", unsafe_allow_html=True)
+    chosen_checkpoint = st.selectbox("Choose a checkpoint", possible_checkpoint + ["all"])
+    if chosen_checkpoint == "all":
+        ds_ckp = ds
+    else:
+        ds_ckp = ds.filter(
+        lambda exs: [lang == chosen_checkpoint for lang in exs["checkpoint"]], batched=True
+    )
+    possible_langs = ds.unique("lang")  
     chosen_lang = st.selectbox("Choose a lang", possible_langs + ["all"])
     if chosen_lang == "all":
-        ds_lang = ds
+        ds_lang = ds_ckp
     else:
-        ds_lang = ds.filter(
-    lambda exs: [lang == chosen_lang for lang in exs["lang"]], batched=True
-)
+        ds_lang = ds_ckp.filter(
+        lambda exs: [lang == chosen_lang for lang in exs["lang"]], batched=True
+    )
+    
     possible_prompts = ds_lang.unique("prompt")
     chosen_prompt = st.selectbox("Choose a prompt", possible_prompts)
     st.markdown(safe_text(chosen_prompt), unsafe_allow_html=True)