filtering

app.py

@@ -79,8 +79,14 @@ def flag(query, language, num_results, issue_description):
     return ""
 
 
-def format_result(result, highlight_terms, exact_search):
+def format_result(result, highlight_terms, exact_search, datasets_filter=None):
     text, url, docid = result
+    if datasets_filter is not None:
+        datasets_filter = set(datasets_filter)
+        dataset = docid.split("/")[1]
+        if not dataset in datasets_filter:
+            return ""
+
     if exact_search:
         query_start = text.find(highlight_terms)
         query_end = query_start + len(highlight_terms)
@@ -119,7 +125,7 @@ def format_result(result, highlight_terms, exact_search):
 
 
 def format_result_page(
-    language, results, highlight_terms, num_results, exact_search
+    language, results, highlight_terms, num_results, exact_search, datasets_filter=None
 ) -> gr.HTML:
     header_html = ""
 
@@ -141,7 +147,7 @@ def format_result_page(
         results_for_lang_html = ""
         for result in results_for_lang:
             results_for_lang_html += format_result(
-                result, highlight_terms, exact_search
+                result, highlight_terms, exact_search, datasets_filter
             )
         if language == "all" and not exact_search:
             results_for_lang_html = f"""
@@ -160,6 +166,7 @@ def extract_results_from_payload(query, language, payload, exact_search):
     results = payload["results"]
 
     processed_results = dict()
+    datasets = set()
     highlight_terms = None
     num_results = None
 
@@ -185,9 +192,11 @@ def extract_results_from_payload(query, language, payload, exact_search):
                 else None
             )
             docid = result["docid"]
+            _, dataset, _ = docid.split("/")
+            datasets.add(dataset)
             processed_results[lang].append((text, url, docid))
 
-    return processed_results, highlight_terms, num_results
+    return processed_results, highlight_terms, num_results, list(datasets)
 
 
 def process_error(error_type):
@@ -238,6 +247,12 @@ if __name__ == "__main__":
     )
 
     with demo:
+        processed_results_state = gr.State([])
+        highlight_terms_state = gr.State([])
+        num_results_state = gr.State(0)
+        exact_search_state = gr.State(False)
+        lang_state = gr.State("")
+
         with gr.Row():
             gr.Markdown(value=description)
         with gr.Row():
@@ -272,20 +287,13 @@ if __name__ == "__main__":
         with gr.Row():
             k = gr.Slider(1, 100, value=10, step=1, label="Max Results")
         with gr.Row():
-            """
-            with gr.Column(scale=1):
-                exact_search = gr.Checkbox(
-                    value=False, label="Exact Search", variant="compact"
-                )
-            """
-            with gr.Column(scale=4):
-                submit_btn = gr.Button("Submit")
+            submit_btn = gr.Button("Submit")
         with gr.Row(visible=False) as datasets_filter:
             available_datasets = gr.Dropdown(
                 type="value",
                 choices=[],
                 value=[],
-                label="Datasets",
+                label="Datasets Filter",
                 multiselect=True,
             )
         with gr.Row():
@@ -310,25 +318,29 @@ if __name__ == "__main__":
             else:
                 query = " ".join(query.split())
             if query == "" or query is None:
-                return "", ""
+                return None
 
             results_html = ""
             payload = request_payload(query, lang, exact_search, k)
             err = extract_error_from_payload(payload)
             if err is not None:
-                results_html = process_error(err)
-            else:
-                (
-                    processed_results,
-                    highlight_terms,
-                    num_results,
-                ) = extract_results_from_payload(query, lang, payload, exact_search)
-                results_html = format_result_page(
-                    lang, processed_results, highlight_terms, num_results, exact_search
-                )
-            datasets = []
-            print(datasets)
+                return process_error(err)
+
+            (
+                processed_results,
+                highlight_terms,
+                num_results,
+                datasets,
+            ) = extract_results_from_payload(query, lang, payload, exact_search)
+            results_html = format_result_page(
+                lang, processed_results, highlight_terms, num_results, exact_search
+            )
+
             return {
+                processed_results_state: processed_results,
+                highlight_terms_state: highlight_terms,
+                num_results_state: num_results,
+                exact_search_state: exact_search,
                 results: results_html,
                 flagging_form: gr.update(visible=True),
                 datasets_filter: gr.update(visible=True),
@@ -337,19 +349,63 @@ if __name__ == "__main__":
                 ),
             }
 
-        def filter_datasets():
-            pass
+        def filter_datasets(
+            lang,
+            processed_results,
+            highlight_terms,
+            num_results,
+            exact_search,
+            datasets_filter,
+        ):
+            results_html = format_result_page(
+                lang,
+                processed_results,
+                highlight_terms,
+                num_results,
+                exact_search,
+                datasets_filter,
+            )
+            return {results: results_html}
 
         query.submit(
             fn=submit,
             inputs=[query, lang, k, available_datasets],
-            outputs=[results, flagging_form, datasets_filter, available_datasets],
+            outputs=[
+                processed_results_state,
+                highlight_terms_state,
+                num_results_state,
+                exact_search_state,
+                results,
+                flagging_form,
+                datasets_filter,
+                available_datasets,
+            ],
         )
         submit_btn.click(
             submit,
             inputs=[query, lang, k, available_datasets],
-            outputs=[results, flagging_form, datasets_filter, available_datasets],
+            outputs=[
+                processed_results_state,
+                highlight_terms_state,
+                num_results_state,
+                exact_search_state,
+                results,
+                flagging_form,
+                datasets_filter,
+                available_datasets,
+            ],
         )
 
-        available_datasets.change(filter_datasets, inputs=[], outputs=[])
+        available_datasets.change(
+            filter_datasets,
+            inputs=[
+                lang,
+                processed_results_state,
+                highlight_terms_state,
+                num_results_state,
+                exact_search_state,
+                available_datasets,
+            ],
+            outputs=[results],
+        )
     demo.launch(enable_queue=True, debug=True)