diff --git a/function.py b/function.py
index 9d768a6..8c1ba8b 100644
--- a/function.py
+++ b/function.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
 from sklearn.preprocessing import StandardScaler
+from sklearn.model_selection import train_test_split
 
 
 id = "202210423130"
@@ -89,13 +90,12 @@ def preprocess_data(X):
 
 
 if __name__ == "__main__":
-    # 加载训练集
-    X_train, y_train = load_data("./data/horse-colic.data")
-    X_train = preprocess_data(X_train)
+    # 加载数据集
+    X, y = load_data("./data/horse-colic.data")
+    X = preprocess_data(X)
 
-    # 加载测试集
-    X_test, y_test = load_data("./data/horse-colic.test")
-    X_test = preprocess_data(X_test)
+    # 使用 train_test_split 打乱并随机抽取数据
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 
     # 训练对率回归模型
     weights = train_logreg(X_train, y_train, learning_rate=0.1, iterations=1000, lambda_reg=0.01)