Sample words based onm frequency

app.py

@@ -47,17 +47,25 @@ def load_vocabulary():
                     token=os.getenv("HF_TOKEN"),
                     filename="vocabulary.csv",
                     local_dir="assets")
-    return pl.read_csv("assets/vocabulary.csv").to_dicts()
+    return pl.read_csv("assets/vocabulary.csv").sort("label_idx")
+
+
+def compute_word_weights(vocabulary):
+    train_counts = vocabulary.group_by("word").agg(pl.col("train_count").sum())
+    total_train_count = train_counts["train_count"].sum()
+    word_weights = [(vocab["word"], vocab["train_count"] / total_train_count)
+                    for vocab in vocabulary.rows(named=True)]
+    words = [word for word, _ in word_weights]
+    weights = [weight for _, weight in word_weights]
+    return words, weights
 
 
 vocabulary = load_vocabulary()
-idx2vocab = {row["word_idx"]: row for row in vocabulary}
-vocab_list = [row["word"] for row in vocabulary]
+words, weights = compute_word_weights(vocabulary)
 
 
-# Select a random word
 def get_random_word():
-    return random.choice(vocab_list)
+    return random.choices(words, weights=weights)[0]
 
 
 # Process the image drawn on canvas
@@ -73,7 +81,7 @@ def process_image(image, current_word):
 
     predictions = []
     for pred, idx in zip(preds_i, indices_i):
-        vocab = idx2vocab[idx]
+        vocab = vocabulary.row(idx, named=True)
         predictions.append({
             "word": vocab["word"],
             "category": vocab["category_name"],
@@ -110,9 +118,6 @@ def create_initial_image():
     return Image.fromarray(data)
 
 
-# Create a white image with the dimensions for the ImageEditor
-initial_image = create_initial_image
-
 # Interface definition
 with gr.Blocks(theme=gr.themes.Soft(),
                css="input {font-size: 24px; font-weight: 600;}") as demo_app:
@@ -121,7 +126,8 @@ with gr.Blocks(theme=gr.themes.Soft(),
 
     with gr.Row():
         word_output = gr.Textbox(label="Your word to draw:",
-                                 value=get_random_word(),
+                                 value=get_random_word,
+                                 interactive=False,
                                  scale=1,
                                  max_lines=1)
         new_word_button = gr.Button("New Word", scale=0, variant="primary")
@@ -133,7 +139,7 @@ with gr.Blocks(theme=gr.themes.Soft(),
             sources=[],
             transforms=[],
             layers=False,
-            value=initial_image,
+            value=create_initial_image,
             brush=gr.Brush(colors=["#000000", "#FF0000", "#00FF00", "#0000FF"],
                            default_size=10))
         plot_output = gr.Plot(label="Model Guesses")