如何按列值将Pandas DataFrame转换为Hugging Face Dataset？

我有如下DataFrame df：

import pandas as pd
from datasets import Dataset

data = [[1, 'Jack', 'A'], [1, 'Jamie', 'A'], [1, 'Mo', 'B'], [1, 'Tammy', 'A'], [2, 'JJ', 'A'], [2, 'Perry', 'C']]
df = pd.DataFrame(data, columns=['id', 'name', 'class'])

输出的df内容：

id   name class
0   1   Jack     A
1   1  Jamie     A
2   1     Mo     B
3   1  Tammy     A
4   2     JJ     A
5   2  Perry     C

我需要将其转换为按id分组、仅含2行的Dataset对象，期望输出如下：

> myDataset
Dataset({
    features: ['id', 'name', 'class'],
    num_rows: 2
})

其中数据结构为：

> myDataset[0:2]
{'id': ['1', '2'], 'name': [['Jack', 'Jamie', 'Mo', 'Tammy'],['JJ', 'Perry']], 'class': [['A', 'A', 'B', 'A'], ['A', 'C']]}

尝试直接用Dataset.from_pandas(df)转换，得到的是6行的Dataset，不符合需求：

myDataset = Dataset.from_pandas(df) 

输出结果：

> myDataset
Dataset({
    features: ['id', 'name', 'class'],
    num_rows: 6
})
> myDataset[0:2]
{'id': [1, 1], 'name': ['Jack', 'Jamie'], 'class': ['A', 'A']}

要实现按id分组的Dataset，需先对DataFrame做分组聚合，再转换为Dataset：

1. 分组聚合DataFrame
   按id分组，将name和class列聚合为列表，id列保留唯一值：

grouped_df = df.groupby('id').agg({
    'name': list,
    'class': list
}).reset_index()

分组后的grouped_df内容：

id                     name               class
0   1  [Jack, Jamie, Mo, Tammy]  [A, A, B, A]
1   2           [JJ, Perry]        [A, C]

2. 转换为Dataset
   将聚合后的DataFrame传入Dataset.from_pandas()：

myDataset = Dataset.from_pandas(grouped_df)

验证结果：

> myDataset
Dataset({
    features: ['id', 'name', 'class'],
    num_rows: 2
})

> myDataset[0:2]
{'id': [1, 2], 'name': [['Jack', 'Jamie', 'Mo', 'Tammy'], ['JJ', 'Perry']], 'class': [['A', 'A', 'B', 'A'], ['A', 'C']]}

完全符合预期需求。

内容的提问来源于stack exchange，提问作者Adrian