Sampling multiple rows from seed dataset for a single record?

[whybe-choi]

Hello, NeMo Data Designer Team!

I'm exploring Data Designer for a use case where I need to sample multiple rows from a seed dataset and use them together as context for generating a single record. For example, I'd like to randomly pick 5 rows from my seed dataset and pass them all to an LLM prompt to generate a multi-document summary. Currently, it seems like the seed dataset works with a 1:1 mapping (one seed row = one generated record). Is there a way to achieve this, or is this something being considered for future releases?

Thanks!

[eric-tramel]

Hiya @whybe-choi :) It's a good question, we've heard this request a few times, it makes sense to find a way to approach this. To help us out a bit, could you show a short example of how you'd like to use this functionality that we can think about in this thread as a kind of psuedo-code MWE? Would something like the below address your use-case?

from data_designer.essentials import (
    DataDesigner, 
    DataDesignerConfigBuilder,
    LLMTextColumnConfig
)

data_designer = DataDesigner()
config_builder = DataDesignerConfigBuilder()



config_builder.with_seed_dataset(
    data_designer.make_seed_reference_from_file("my_data.csv"),
    sampling_strategy=SamplingStrategy.SHUFFLED,
   
    # New option here?
    take=5
)

This option could do something like the opposite of an .explode() on the dataset, and would nest lists of potentially ordered or shuffled entries. Assuming that we have a field called "text" within the original dataset, we go from dataset["text"] = "some text" to a format like dataset["text"] = ["text 1", "text 42", "text34", "text12", "text99"]. These could maybe be iterated over from the Jinja template in a prompt.

config_builder.add_column(LLMTextColumnConfig(prompt="""\
{% for doc in text %}
<document>
{{ doc }}
</document>
{% endfor %}

Compare and contrast these documents.\
"""))

[whybe-choi]

Hi, @eric-tramel !

That’s exactly the feature I was looking for! My use case involves processing PDF images—I’ve parsed them into Markdown and stored them in the seed dataset. I want to provide multiple Markdown entries (pages) to the LLM as a single context to generate summaries across several pages. Your proposed pseudo-code would address this perfectly.