add "see more" button

app.py

@@ -1,6 +1,7 @@
 import time
+from itertools import count, islice
 from functools import partial
-from typing import Iterator
+from typing import Iterable, Iterator, TypeVar
 
 import gradio as gr
 import requests.exceptions
@@ -10,9 +11,12 @@ from huggingface_hub import InferenceClient
 model_id = "microsoft/Phi-3-mini-4k-instruct"
 client = InferenceClient(model_id)
 
+MAX_TOTAL_NB_ITEMS = 100  # almost infinite, don't judge me (actually it's because gradio needs a fixed number of components)
+MAX_NB_ITEMS_PER_GENERATION_CALL = 10
+
 GENERATE_DATASET_NAMES_FOR_SEARCH_QUERY = (
         "A Machine Learning Practioner is looking for a dataset that matches '{search_query}'. "
-        "Generate a list of 10 names of quality dataset that don't exist but sound plausible and would "
+        f"Generate a list of {MAX_NB_ITEMS_PER_GENERATION_CALL} names of quality dataset that don't exist but sound plausible and would "
         "be helpful. Feel free to reuse words from the query '{search_query}' to name the datasets. "
         "Every dataset should be about '{search_query}' and have descriptive tags/keywords including the ML task name associated to the dataset (classification, regression, anomaly detection, etc.). Use the following format:\n1. DatasetName1 (tag1, tag2, tag3)\n1. DatasetName2 (tag1, tag2, tag3)"
 )
@@ -25,52 +29,9 @@ GENERATE_DATASET_CONTENT_FOR_SEARCH_QUERY_AND_NAME_AND_TAGS = (
     "Reply using a short description of the dataset with title **Dataset Description:** followed by the CSV content in a code block and with title **CSV Content Preview:**."
 )
 
-default_query = "various datasets on many different subjects and topics, from classification to language modeling, from science to sport to finance to news"
-
-
-def stream_reponse(msg: str, max_tokens=500) -> Iterator[str]:
-    for _ in range(3):
-        try:
-            for message in client.chat_completion(
-                messages=[{"role": "user", "content": msg}],
-                max_tokens=max_tokens,
-                stream=True,
-            ):
-                yield message.choices[0].delta.content
-        except requests.exceptions.ConnectionError as e:
-            print(e + "\n\nRetrying in 1sec")
-            time.sleep(1)
-            continue
-        break
-
-
-def gen_datasets(search_query: str) -> Iterator[str]:
-    search_query = search_query[:1000] if search_query.strip() else default_query
-    generated_text = ""
-    for token in stream_reponse(GENERATE_DATASET_NAMES_FOR_SEARCH_QUERY.format(search_query=search_query)):
-        generated_text += token
-        if generated_text.endswith("\n"):
-            yield generated_text.strip()
-    yield generated_text.strip()
-    print("-----\n\n" + generated_text)
-
-
-def gen_dataset_content(search_query: str, dataset_name: str, tags: str) -> Iterator[str]:
-    search_query = search_query[:1000] if search_query.strip() else default_query
-    generated_text = ""
-    for token in stream_reponse(GENERATE_DATASET_CONTENT_FOR_SEARCH_QUERY_AND_NAME_AND_TAGS.format(
-        search_query=search_query,
-        dataset_name=dataset_name,
-        tags=tags,
-    ), max_tokens=1500):
-        generated_text += token
-        yield generated_text
-    print("-----\n\n" + generated_text)
-
-
-NB_ITEMS_PER_PAGE = 10
-
-default_output = """
+landing_page_query = "various datasets on many different subjects and topics, from classification to language modeling, from science to sport to finance to news"
+
+landing_page_datasets_generated_text = """
 1. NewsEventsPredict (classification, media, trend)
 2. FinancialForecast (economy, stocks, regression)
 3. HealthMonitor (science, real-time, anomaly detection)
@@ -81,10 +42,15 @@ default_output = """
 8. NewsEventTracker (classification, public awareness, topical clustering)
 9. HealthVitalSigns (anomaly detection, biometrics, prediction)
 10. GameStockPredict (classification, finance, sports contingency)
-""".strip().split("\n")
-assert len(default_output) == NB_ITEMS_PER_PAGE
+"""
+default_output = landing_page_datasets_generated_text.strip().split("\n")
+assert len(default_output) == MAX_NB_ITEMS_PER_GENERATION_CALL
 
 css = """
+a {
+    color: var(--body-text-color);
+}
+
 .datasetButton {
     justify-content: start;
     justify-content: left;
@@ -93,9 +59,6 @@ css = """
     font-size: var(--button-small-text-size);
     color: var(--body-text-color-subdued);
 }
-a {
-    color: var(--body-text-color);
-}
 .topButton {
     justify-content: start;
     justify-content: left;
@@ -134,6 +97,10 @@ a {
 .buttonsGroup div {
     background: transparent;
 }
+.insivibleButtonGroup {
+    display: none;
+}
+
 @keyframes placeHolderShimmer{
     0%{
         background-position: -468px 0
@@ -155,39 +122,9 @@ a {
 }
 """
 
-def search_datasets(search_query):
-    output_values = [
-        gr.Button("⬜⬜⬜⬜⬜⬜", elem_classes="topButton linear-background"),
-        gr.Button("░░░░, ░░░░, ░░░░", elem_classes="bottomButton linear-background")
-    ] * NB_ITEMS_PER_PAGE
-    for generated_text in gen_datasets(search_query):
-        if "I'm sorry" in generated_text:
-            raise gr.Error("Error: inappropriate content")
-        lines = [line for line in generated_text.split("\n") if line and line.split(".", 1)[0].isnumeric()][:NB_ITEMS_PER_PAGE]
-        for i, line in enumerate(lines):
-            dataset_name, tags = line.split(".", 1)[1].strip(" )").split(" (", 1)
-            output_values[2 * i] = gr.Button(dataset_name, elem_classes="topButton")
-            output_values[2 * i + 1] = gr.Button(tags, elem_classes="bottomButton")
-        yield output_values
-
-
-def show_dataset(search_query, *buttons_values, i):
-    dataset_name, tags = buttons_values[2 * i : 2 * i + 2]
-    dataset_title = f"# {dataset_name}\n\n tags: {tags}\n\n _Note: This is an AI-generated dataset so its content may be inaccurate or false_"
-    yield gr.Column(visible=False), gr.Column(visible=True), dataset_title, ""
-    for generated_text in gen_dataset_content(search_query=search_query, dataset_name=dataset_name, tags=tags):
-        yield gr.Column(), gr.Column(), dataset_title, generated_text
-
-
-def show_search_page():
-    return gr.Column(visible=True), gr.Column(visible=False)
-
-
-def generate_full_dataset():
-    raise gr.Error("Not implemented yet sorry ! Give me some feedbacks in the Community tab in the meantime ;)")
-
 
 with gr.Blocks(css=css) as demo:
+    generated_texts_state = gr.State((landing_page_datasets_generated_text,))
     with gr.Row():
         with gr.Column(scale=4, min_width=0):
             pass
@@ -208,28 +145,34 @@ with gr.Blocks(css=css) as demo:
                 search_button = gr.Button("🔍", variant="primary")
             with gr.Column(scale=4, min_width=0):
                 pass
-        inputs = [search_bar]
-        show_dataset_outputs = [search_page]
         with gr.Row():
             with gr.Column(scale=4, min_width=0):
                 pass
             with gr.Column(scale=10):
-                buttons = []
-                for i in range(10):
-                    line = default_output[i]
-                    dataset_name, tags = line.split(".", 1)[1].strip(" )").split(" (", 1)
-                    with gr.Group(elem_classes="buttonsGroup"):
-                        top = gr.Button(dataset_name, elem_classes="topButton")
-                        bottom = gr.Button(tags, elem_classes="bottomButton")
-                        buttons += [top, bottom]
-                        top.click(partial(show_dataset, i=i), inputs=inputs, outputs=show_dataset_outputs)
-                        bottom.click(partial(show_dataset, i=i), inputs=inputs, outputs=show_dataset_outputs)
-                inputs += buttons
+                button_groups: list[gr.Group] = []
+                buttons: list[gr.Button] = []
+                for i in range(MAX_TOTAL_NB_ITEMS):
+                    if i < len(default_output):
+                        line = default_output[i]
+                        dataset_name, tags = line.split(".", 1)[1].strip(" )").split(" (", 1)
+                        group_classes = "buttonsGroup"
+                        dataset_name_classes = "topButton"
+                        tags_classes = "bottomButton"
+                    else:
+                        dataset_name, tags = "⬜⬜⬜⬜⬜⬜", "░░░░, ░░░░, ░░░░"
+                        group_classes = "buttonsGroup insivibleButtonGroup"
+                        dataset_name_classes = "topButton linear-background"
+                        tags_classes = "bottomButton linear-background"
+                    with gr.Group(elem_classes=group_classes) as button_group:
+                        button_groups.append(button_group)
+                        buttons.append(gr.Button(dataset_name, elem_classes=dataset_name_classes))
+                        buttons.append(gr.Button(tags, elem_classes=tags_classes))
+
+                see_more = gr.Button("See more")  # TODO: dosable when reaching end of page
                 gr.Markdown(f"_powered by [{model_id}](https://huggingface.co/{model_id})_")
             with gr.Column(scale=4, min_width=0):
                 pass
-        search_bar.submit(search_datasets, inputs=search_bar, outputs=buttons)
-        search_button.click(search_datasets, inputs=search_bar, outputs=buttons)
+        # more.click(search_more_datasets, inputs=[generated_texts, search_bar], outputs=[generated_texts] + buttons)
     with gr.Column(visible=False) as dataset_page:
         with gr.Row():
             with gr.Column(scale=4, min_width=0):
@@ -241,13 +184,176 @@ with gr.Blocks(css=css) as demo:
                     with gr.Column(scale=4, min_width=0):
                         pass
                     with gr.Column():
-                        generate_full_dataset_button = gr.Button("Generate Full Dataset", variant="primary")
-                        generate_full_dataset_button.click(generate_full_dataset)
+                        generate_full_dataset_button = gr.Button("Generate Full Dataset", variant="primary")  # TODO: implement
                         back_button = gr.Button("< Back", size="sm")
-                        back_button.click(show_search_page, inputs=[], outputs=[search_page, dataset_page])
                     with gr.Column(scale=4, min_width=0):
                         pass
             with gr.Column(scale=4, min_width=0):
                 pass
-        show_dataset_outputs += [dataset_page, dataset_title, dataset_content]
+
+
+    T = TypeVar("T")
+
+    def batched(it: Iterable[T], n: int) -> Iterator[list[T]]:
+        it = iter(it)
+        while batch := list(islice(it, n)):
+            yield batch
+
+
+    def stream_reponse(msg: str, generated_texts: tuple[str] = (), max_tokens=500) -> Iterator[str]:
+        messages = [
+            {"role": "user", "content": msg}
+        ] + [
+            item
+            for generated_text in generated_texts
+            for item in [
+                {"role": "assistant", "content": generated_text},
+                {"role": "user", "content": "Can you generate more ?"},
+            ]
+        ]
+        for _ in range(3):
+            try:
+                for message in client.chat_completion(
+                    messages=messages,
+                    max_tokens=max_tokens,
+                    stream=True,
+                    top_p=0.8,
+                ):
+                    yield message.choices[0].delta.content
+            except requests.exceptions.ConnectionError as e:
+                print(e + "\n\nRetrying in 1sec")
+                time.sleep(1)
+                continue
+            break
+
+
+    def gen_datasets_line_by_line(search_query: str, generated_texts: tuple[str] = ()) -> Iterator[str]:
+        search_query = search_query[:1000] if search_query.strip() else landing_page_query
+        generated_text = ""
+        current_line = ""
+        for token in stream_reponse(
+            GENERATE_DATASET_NAMES_FOR_SEARCH_QUERY.format(search_query=search_query),
+            generated_texts=generated_texts,
+        ):
+            current_line += token
+            if current_line.endswith("\n"):
+                yield current_line
+                generated_text += current_line
+                current_line = ""
+        yield current_line
+        generated_text += current_line
+        print("-----\n\n" + generated_text)
+
+
+    def gen_dataset_content(search_query: str, dataset_name: str, tags: str) -> Iterator[str]:
+        search_query = search_query[:1000] if search_query.strip() else landing_page_query
+        generated_text = ""
+        for token in stream_reponse(GENERATE_DATASET_CONTENT_FOR_SEARCH_QUERY_AND_NAME_AND_TAGS.format(
+            search_query=search_query,
+            dataset_name=dataset_name,
+            tags=tags,
+        ), max_tokens=1500):
+            generated_text += token
+            yield generated_text
+        print("-----\n\n" + generated_text)
+
+    search_datasets_inputs = search_bar
+    search_datasets_outputs = button_groups + buttons + [generated_texts_state]
+
+    def search_datasets(search_query):
+        yield {generated_texts_state: []}
+        yield {
+            button_group: gr.Group(elem_classes="buttonsGroup insivibleButtonGroup")
+            for button_group in button_groups[MAX_NB_ITEMS_PER_GENERATION_CALL:]
+        }
+        yield {
+            k: v
+            for dataset_name_button, tags_button in batched(buttons, 2)
+            for k, v in {
+                dataset_name_button: gr.Button("⬜⬜⬜⬜⬜⬜", elem_classes="topButton linear-background"),
+                tags_button: gr.Button("░░░░, ░░░░, ░░░░", elem_classes="bottomButton linear-background")
+            }.items()
+        }
+        current_item_idx = 0
+        generated_text = ""
+        for line in gen_datasets_line_by_line(search_query):
+            if "I'm sorry" in line:
+                raise gr.Error("Error: inappropriate content")
+            if current_item_idx >= MAX_NB_ITEMS_PER_GENERATION_CALL:
+                return
+            if line.strip() and line.strip().split(".", 1)[0].isnumeric():
+                try:
+                    dataset_name, tags = line.strip().split(".", 1)[1].strip(" )").split(" (", 1)
+                except ValueError:
+                    dataset_name, tags = line.strip().split(".", 1)[1].strip(" )").split(" ", 1)
+                dataset_name, tags = dataset_name.strip("()[]* "), tags.strip("()[]* ")
+                generated_text += line
+                yield {
+                    buttons[2 * current_item_idx]: gr.Button(dataset_name, elem_classes="topButton"),
+                    buttons[2 * current_item_idx + 1]: gr.Button(tags, elem_classes="bottomButton"),
+                    generated_texts_state: (generated_text,),
+                }
+                current_item_idx += 1
+
+    search_more_datasets_inputs = [search_bar, generated_texts_state]
+    search_more_datasets_outputs = button_groups + buttons + [generated_texts_state]
+
+    def search_more_datasets(search_query, generated_texts):
+        current_item_idx = initial_item_idx = len(generated_texts) * MAX_NB_ITEMS_PER_GENERATION_CALL
+        yield {
+            button_group: gr.Group(elem_classes="buttonsGroup")
+            for button_group in button_groups[len(generated_texts) * MAX_NB_ITEMS_PER_GENERATION_CALL:(len(generated_texts) + 1) * MAX_NB_ITEMS_PER_GENERATION_CALL]
+        }
+        generated_text = ""
+        for line in gen_datasets_line_by_line(search_query, generated_texts=generated_texts):
+            if "I'm sorry" in line:
+                raise gr.Error("Error: inappropriate content")
+            if current_item_idx - initial_item_idx >= MAX_NB_ITEMS_PER_GENERATION_CALL:
+                return
+            if line.strip() and line.strip().split(".", 1)[0].isnumeric():
+                try:
+                    dataset_name, tags = line.strip().split(".", 1)[1].strip(" )").split(" (", 1)
+                except ValueError:
+                    dataset_name, tags = line.strip().split(".", 1)[1].strip(" )").split(" ", 1) [0], ""
+                dataset_name, tags = dataset_name.strip("()[]* "), tags.strip("()[]* ")
+                generated_text += line
+                yield {
+                    buttons[2 * current_item_idx]: gr.Button(dataset_name, elem_classes="topButton"),
+                    buttons[2 * current_item_idx + 1]: gr.Button(tags, elem_classes="bottomButton"),
+                    generated_texts_state: (*generated_texts, generated_text),
+                }
+                current_item_idx += 1
+
+    show_dataset_inputs = [search_bar, *buttons]
+    show_dataset_outputs = [search_page, dataset_page, dataset_title, dataset_content]
+
+    def show_dataset(search_query, *buttons_values, i):
+        dataset_name, tags = buttons_values[2 * i : 2 * i + 2]
+        yield {
+            search_page: gr.Column(visible=False),
+            dataset_page: gr.Column(visible=True),
+            dataset_title: f"# {dataset_name}\n\n tags: {tags}\n\n _Note: This is an AI-generated dataset so its content may be inaccurate or false_"
+        }
+        for generated_text in gen_dataset_content(search_query=search_query, dataset_name=dataset_name, tags=tags):
+            yield {dataset_content: generated_text}
+
+
+    def show_search_page():
+        return gr.Column(visible=True), gr.Column(visible=False)
+
+
+    def generate_full_dataset():
+        raise gr.Error("Not implemented yet sorry ! Give me some feedbacks in the Community tab in the meantime ;)")
+
+
+    search_bar.submit(search_datasets, inputs=search_datasets_inputs, outputs=search_datasets_outputs)
+    search_button.click(search_datasets, inputs=search_datasets_inputs, outputs=search_datasets_outputs)
+    for i, (dataset_name_button, tags_button) in enumerate(batched(buttons, 2)):
+        dataset_name_button.click(partial(show_dataset, i=i), inputs=show_dataset_inputs, outputs=show_dataset_outputs)
+        tags_button.click(partial(show_dataset, i=i), inputs=show_dataset_inputs, outputs=show_dataset_outputs)
+    see_more.click(search_more_datasets, inputs=search_more_datasets_inputs, outputs=search_more_datasets_outputs)
+
+    generate_full_dataset_button.click(generate_full_dataset)
+    back_button.click(show_search_page, inputs=[], outputs=[search_page, dataset_page])
+
 demo.launch()