saulpw / aipl

Array-Inspired Pipeline Language
MIT License
119 stars 7 forks source link

make all combinations of elements in multiple columns #34

Closed dovinmu closed 1 year ago

dovinmu commented 1 year ago

It might be helpful to have the cartesian product of two or more columns. The use case here is iterating over the combination of the set of model IDs and tasks for BIG bench. Here's a toy example that I got working. I used the csv-load op in #30 on the following CSV:

agent,problem
pikachu,needing to recharge
poseidon,being too wet
!!python
from aipl.table import Table

@defop('cartesian-product', 1.5, 1.5)
def _cartesian(aipl, t:Table, col1, col2) -> dict:
    column1 = []
    column2 = []
    for row in t:
        column1.append(row[col1])
        column2.append(row[col2])
    ret = []
    for el1 in column1:
        for el2 in column2:
            ret.append({col1: el1, col2: el2})
    return ret

!csv-parse test.csv
!cartesian-product agent problem

!format
Please respond with a succinct suggestion for {agent} struggling with {problem}.
!llm
!format
Q: Please respond with a succinct suggestion for {agent} struggling with {problem}.
A: {_}
---
!print

Ideally we wouldn't need to specify the column names.