add comments

app.py

@@ -1,9 +1,19 @@
-from datasets import load_dataset
 import streamlit as st
+from datasets import load_dataset
+
+st.set_page_config(
+    page_icon="🧊",
+    layout="wide",
+)
+
+st.write(
+    "This is an application for viewing different generations for the same prompt. The generations vary depending on the checkpoint used and also the parameters used for the generation."
+)
 
 HF_API_TOKEN = st.secrets["HF_API_TOKEN"]
 PROMPT_COLOR = "#CA437E"
 
+
 def safe_text(text):
     text = text.replace("\n", "<br>")
     return f"<pre>{text}</pre>"
@@ -16,20 +26,44 @@ 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 = 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)
+col_1, col_2 = st.columns(2)
+with col_1:
+    st.markdown("<h1 style='text-align: center'>Prompt</h1>", unsafe_allow_html=True)
+    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(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
+with col_2:
+    st.markdown(
+        "<h1 style='text-align: center'>Generation</h1>", unsafe_allow_html=True
+    )
+    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)
+    st.markdown(
+        "<h2 style='text-align: center'>Generation configuration</h2>",
+        unsafe_allow_html=True,
+    )
+    config = {
+        key: value
+        for key, value in sample.items()
+        if key not in ["prompt", "generation"]
+    }
+    config