Generate multiple alternatives for image search

app.py

@@ -23,6 +23,9 @@ if "report_submitted" not in st.session_state:
 if "grascii" not in st.session_state:
     st.session_state["grascii"] = ""
 
+if "alternatives" not in st.session_state:
+    st.session_state["alternatives"] = {}
+
 if st.session_state["report_submitted"]:
     st.toast("Thanks for the report!")
     st.session_state["report_submitted"] = False

search.py

@@ -31,9 +31,10 @@ def load_images():
 image_map = load_images()
 
 
-def set_grascii():
+def on_submit():
     if "grascii_text_box" in st.session_state:
         st.session_state["grascii"] = st.session_state["grascii_text_box"]
+        st.session_state["alternatives"] = {}
 
 
 def write_grascii_search():
@@ -46,7 +47,10 @@ def write_grascii_search():
         placeholder = st.empty()
         if search_by == "text":
             placeholder.text_input(
-                "Grascii", value=st.session_state["grascii"], key="grascii_text_box", max_chars=MAX_GRASCII_LENGTH
+                "Grascii",
+                value=st.session_state["grascii"],
+                key="grascii_text_box",
+                max_chars=MAX_GRASCII_LENGTH,
             )
         else:
             image_data = placeholder.file_uploader(
@@ -74,10 +78,14 @@ def write_grascii_search():
                 alpha_composite = Image.alpha_composite(background, image)
 
                 arr = np.array([alpha_composite.convert("L")])
-                tokens = run_vision(arr)
-                st.session_state["grascii"] = "".join(tokens)
+                predictions = run_vision(arr)
+                alternatives = {"".join(p): True for p in predictions}
+                if st.session_state["alternatives"] != alternatives:
+                    st.session_state["alternatives"] = alternatives
+                    st.session_state["grascii"] = "".join(predictions[0])
+
                 if save:
-                    save_image(image_data.getvalue(), "-".join(tokens))
+                    save_image(image_data.getvalue(), "-".join(predictions[0]))
 
         with st.expander("Options"):
             interpretation = st.radio(
@@ -157,7 +165,7 @@ def write_grascii_search():
                     """,
             )
 
-        st.form_submit_button("Search", on_click=set_grascii)
+        st.form_submit_button("Search", on_click=on_submit)
 
     grascii = st.session_state["grascii"]
 
@@ -180,9 +188,24 @@ def write_grascii_search():
         if grascii:
             st.error(f"Invalid Grascii\n```\n{e.context}\n```")
     else:
+        if len(st.session_state["alternatives"]) > 1:
+            st.pills(
+                "Alternatives",
+                st.session_state["alternatives"],
+                key="alternative",
+                default=grascii,
+                on_change=on_alternative_selection,
+            )
         write_results(grascii_results, grascii.upper(), "grascii")
 
 
+def on_alternative_selection():
+    if st.session_state["alternative"] is None:
+        st.session_state["alternative"] = st.session_state["grascii"]
+    else:
+        st.session_state["grascii"] = st.session_state["alternative"]
+
+
 @st.fragment
 def write_results(results, term, key_prefix):
     rows = map(

vision.py

@@ -1,3 +1,5 @@
+import math
+
 import streamlit as st
 from transformers import (
     PreTrainedTokenizerFast,
@@ -5,25 +7,45 @@ from transformers import (
     ViTImageProcessor,
 )
 
-model_name = "grascii/gregg-vision-v0.2.1"
+
+MODEL_NAME = "grascii/gregg-vision-v0.2.1"
+MIN_LOG_PROB = math.log(0.5)
+NUM_BEAMS = 3
 
 
-@st.cache_resource(show_spinner=f"Loading {model_name}")
+@st.cache_resource(show_spinner=f"Loading {MODEL_NAME}")
 def load_model():
     model = VisionEncoderDecoderModel.from_pretrained(
-        model_name, token=st.secrets.HF_TOKEN
+        MODEL_NAME, token=st.secrets.HF_TOKEN
     )
     tokenizer = PreTrainedTokenizerFast.from_pretrained(
-        model_name,
+        MODEL_NAME,
         token=st.secrets.HF_TOKEN,
     )
-    processor = ViTImageProcessor.from_pretrained(model_name, token=st.secrets.HF_TOKEN)
+    processor = ViTImageProcessor.from_pretrained(MODEL_NAME, token=st.secrets.HF_TOKEN)
     return model, tokenizer, processor
 
 
-@st.cache_data(ttl=3600, show_spinner=f"Running {model_name}")
+@st.cache_data(ttl=3600, show_spinner=f"Running {MODEL_NAME}")
 def run_vision(image):
     model, tokenizer, processor = load_model()
     pixel_values = processor(image, return_tensors="pt").pixel_values
-    generated = model.generate(pixel_values, max_new_tokens=12)[0]
-    return tokenizer.convert_ids_to_tokens(generated, skip_special_tokens=True)
+    generated = model.generate(
+        pixel_values,
+        max_new_tokens=12,
+        num_beams=NUM_BEAMS,
+        num_return_sequences=NUM_BEAMS,
+        output_scores=True,
+        return_dict_in_generate=True,
+    )
+    return [
+        tokenizer.convert_ids_to_tokens(
+            generated["sequences"][0], skip_special_tokens=True
+        )
+    ] + [
+        tokenizer.convert_ids_to_tokens(seq, skip_special_tokens=True)
+        for seq, score in zip(
+            generated["sequences"][1:], generated["sequences_scores"][1:]
+        )
+        if score > MIN_LOG_PROB
+    ]