如何使用IFELSE为具有多个概率的"PimaIndiansDiabetes“将数据拆分成训练/测试集？

import pandas as pd
from sklearn.model_selection import train_test_split
import numpy as np

# 加载数据集
data = pd.read_csv("PimaIndiansDiabetes.csv")

# 定义概率阈值
threshold = 0.8

# 将数据集拆分为训练集和测试集
train_data = pd.DataFrame()
test_data = pd.DataFrame()

for index, row in data.iterrows():
    # 生成随机数
    random_num = np.random.random()
    
    # 判断随机数是否小于概率阈值
    if random_num < threshold:
        train_data = train_data.append(row)
    else:
        test_data = test_data.append(row)

# 输出训练集和测试集的大小
print("训练集大小：", train_data.shape)
print("测试集大小：", test_data.shape)