update app

app.py

@@ -4,10 +4,6 @@ import streamlit as st
 HF_API_TOKEN = st.secrets["HF_API_TOKEN"]
 PROMPT_COLOR = "#CA437E"
 
-ds = load_dataset("SaulLu/bloom-generations",  use_auth_token=HF_API_TOKEN)
-ds = ds["train"]
-
-
 def safe_text(text):
     text = text.replace("\n", "<br>")
     return f"<pre>{text}</pre>"
@@ -20,10 +16,20 @@ def prompt_markup_format(text):
 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 = ds["train"]
+
+possible_prompts = ds.unique("prompt")
+chosen_prompt = st.selectbox("Chose a prompt", possible_prompts)
+st.markdown(safe_text(chosen_prompt), unsafe_allow_html=True)
+
+sub_ds = ds.filter(lambda exs:[prompt==chosen_prompt for prompt in exs["prompt"]], batched=True)
+
 
-index_sample = st.number_input("Index of the chosen example", min_value=0, max_value=len(ds) - 1, value=0, step=1)
-sample = ds[index_sample]
-markdown_text = f"{prompt_markup_format(sample['prompt'])}{generation_markup_format(sample['generation'])}"
-markdown_text = f"{safe_text(sample['prompt'])}{generation_markup_format(safe_text(sample['generation']))}"
+index_sample = st.number_input("Index of the chosen example", min_value=0, max_value=len(sub_ds) - 1, value=0, step=1)
+sample = sub_ds[index_sample]
+markdown_text = generation_markup_format(safe_text(sample['generation']))
 
 st.markdown(markdown_text, unsafe_allow_html=True)
+config = {key:value for key, value in sample.items() if key not in ["prompt", "generation"]}
+config