regex for groups

app.py

@@ -2,6 +2,7 @@ from concurrent.futures import ThreadPoolExecutor
 import enum
 from functools import partial
 import json
+import os
 from pathlib import Path
 import re
 import tempfile
@@ -26,6 +27,8 @@ LOG_SCALE_STATS = {
     "pages_with_lorem_ipsum",
 }
 
+STATS_LOCATION_DEFAULT = os.getenv("STATS_LOCATION_DEFAULT", "s3://")
+
 
 def find_folders(base_folder, path):
     base_folder = get_datafolder(base_folder)
@@ -145,11 +148,12 @@ def prepare_non_grouped_data(dataset_path, base_folder, grouping, stat_name, nor
     return stats_rounded
 
 
-def prepare_grouped_data(dataset_path, base_folder, grouping, stat_name, top_k, direction: PARTITION_OPTIONS):
+def prepare_grouped_data(dataset_path, base_folder, grouping, stat_name, top_k, direction: PARTITION_OPTIONS, regex):
     import heapq
+    regex_compiled = re.compile(regex) if regex else None
 
     stats = load_stats(base_folder, dataset_path, stat_name, grouping)
-
+    stats = {key: value for key, value in stats.items() if not regex or regex_compiled.match(key)}
     means = {key: value.mean for key, value in stats.items()}
 
     # Use heap to get top_k keys
@@ -254,6 +258,7 @@ def update_graph(
     normalization,
     top_k,
     direction,
+    regex,
     progress=gr.Progress(),
 ):
     if len(datasets) <= 0 or not stat_name or not grouping:
@@ -262,7 +267,7 @@ def update_graph(
     prepare_fc = (
         partial(prepare_non_grouped_data, normalization=normalization)
         if grouping == "histogram"
-        else partial(prepare_grouped_data, top_k=top_k, direction=direction)
+        else partial(prepare_grouped_data, top_k=top_k, direction=direction, regex=regex)
     )
     graph_fc = (
         partial(plot_scatter, normalization=normalization)
@@ -348,20 +353,27 @@ Groupings:
                 )
 
             with gr.Row(visible=False) as group_choices:
-                top_select = gr.Number(
-                    label="K",
-                    value=100,
-                    interactive=True,
-                )
-                direction_checkbox = gr.Radio(
-                    label="Partition",
-                    choices=[
-                        "Top",
-                        "Bottom",
-                        "Most frequent (n_docs)",
-                    ],
-                    value="Top",
-                )
+                with gr.Column(scale=2):
+                    group_regex = gr.Text(
+                        label="Group Regex",
+                        value=None,
+                    )
+                    with gr.Row():
+                        top_select = gr.Number(
+                            label="N Groups",
+                            value=100,
+                            interactive=True,
+                        )
+                    
+                        direction_checkbox = gr.Radio(
+                            label="Partition",
+                            choices=[
+                                "Top",
+                                "Bottom",
+                                "Most frequent (n_docs)",
+                            ],
+                            value="Most frequent (n_docs)",
+                        )
 
             update_button = gr.Button("Update Graph", variant="primary")
     with gr.Row():
@@ -413,6 +425,7 @@ Groupings:
             normalization_checkbox,
             top_select,
             direction_checkbox,
+            group_regex,
         ],
         outputs=[graph_output, exported_data, export_data_button],
     )