add lang

app.py

@@ -27,18 +27,26 @@ def generation_markup_format(text):
     return f"<font color={PROMPT_COLOR}>{text}</pre></font>"
 
 
-ds = load_dataset("SaulLu/bloom-generations", use_auth_token=HF_API_TOKEN)
+ds = load_dataset("/home/lucile_huggingface_co/repos/bloom-generations/bloom-generations.py", use_auth_token=HF_API_TOKEN)
 ds = ds["train"]
 
-possible_prompts = ds.unique("prompt")
+possible_langs = ds.unique("lang")
 
 col_1, col_2 = st.columns(2)
 with col_1:
     st.markdown("<h1 style='text-align: center'>Prompt</h1>", unsafe_allow_html=True)
+    chosen_lang = st.selectbox("Choose a lang", possible_langs + ["all"])
+    if chosen_lang == "all":
+        ds_lang = ds
+    else:
+        ds_lang = ds.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)
 
-sub_ds = ds.filter(
+sub_ds = ds_lang.filter(
     lambda exs: [prompt == chosen_prompt for prompt in exs["prompt"]], batched=True
 )