Add data transformation capability to dataset retrieval step

[saum7800]

Description

We are adding a new component DatasetTransformer. The code currently contains one version of DatasetTransformer: PromptBasedDatasetTransformer. It is used with the Dataset Retrieval step in order to transform retrieved data to a format that is more directly relevant to the task given by a user. Here is the broad flow:

1. Dataset retriever chooses a dataset
2. Automatic column selection chooses relevant columns from the dataset
3. We remove columns from the dataset that were not chosen by the automatic column selection
4. We define a new PromptBasedDatasetTransformer object, and call the transform_data function with the PromptSpec and the loaded dataset.
5. transform_data creates a prompt and calls the APIAgent which returns a "plan" for carrying out the transformation. This prompt uses the prompt_spec and example rows of the retrieved dataset.
6. Next, it creates a list of "transform_prompts", where each transform prompt contains the "plan", one sample from the retrieved dataset, and the PromptSpec. It then requests a batch complete from the APIAgent for all the transform_prompts.
7. Finally, we extract the "input" and "output" keys from each response, and canonicalize the inputs and outputs into a dataset.

[neubig]

Thanks a lot! This looks good but I have some suggested changes.

[neubig]

Hi @saum7800 I can take another look when you've had a moment to make the above revisions!

[saum7800]

Hey @neubig , I have resolved the comments that seemed like easy fixes, and left comments for a couple of them we can discuss. Please re-review whenever you get a chance. Thanks!

[neubig]

Looks good, thanks for making the changes!

[viswavi]

Few changes required, particularly about deduplicating functions that do similar things and improving function/variable names

[viswavi]

This is fine, but it suggests that this function is not well-typed...

[saum7800]

I think I had to do it because certain ones had a train split that had to be accessed to access samples, whereas some had direct access to samples. either way, I suggest we keep this as it is for now and resolve in a future PR, since that is not the focus for the PR. have tested that this works for now

[viswavi]

I feel that now we're conflating two concepts:

• canonicalization (which I thought of as just column selection)
• transformation (which may make major changes to the dataset's contents, thus not really making it more canonical, according to the definition of that word)

[saum7800]

I have merged such that reranking happens separately and colsel+ (if user specifies) transformation separately. we can discuss if any renaming is required in your re-review.

[viswavi]

Finished my first pass here.

It looks like you might need to pull/merge latest main and update your base branch, since some of the files being modified here (e.g. parse_json_responses.py) don't exist on main

[viswavi]

Looks great! I've left a few minor comments, all focusing on superficial things. The general design and code quality is great.

Feel free to merge this after addressing the few remaining comments.