python——在联合学习中将数据分为培训和测试

我是联邦学习方面的新手，目前正在按照TFF官方文档进行模型试验。但我被这个问题困住了，希望能在这里找到一些解释。

当我使用自己的数据集时，数据分布在多个文件中，每个文件都是一个客户机（因为我计划构建模型）。定义了相关变量和自变量。

现在，我的问题是，在联邦学习中，如何将数据拆分为每个客户机（文件）中的训练集和测试集？正如我们通常在集中化的ML模型中所做的，如下面的代码所示

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.4, random_state = 42) 

我一直在找这个邮递这与我的问题类似，我也遵循了那里的说明，但它的目的是将客户划分为培训和测试客户本身，而不是这些客户机中的数据。