diff --git a/function.py b/function.py
index d9b0ba1..fbd7ac4 100644
--- a/function.py
+++ b/function.py
@@ -29,19 +29,22 @@ name = "章程鸿"
 import numpy as np
 from sklearn.preprocessing import StandardScaler
 
-
 def sigmoid(z):
     return 1.0 / (1 + np.exp(-z))
 
+def compute_gradient(X_batch, y_batch, weights):
+    z = np.dot(X_batch, weights)
+    h = sigmoid(z)
+    gradient = np.dot(X_batch.T, (h - y_batch)) / len(y_batch)
+    return gradient
 
 def train_logreg(X, y, learning_rate=0.05, iterations=1500, batch_size=64):
     m, n = X.shape
     weights = np.random.randn(n) * 0.01
     best_weights = weights.copy()
     best_accuracy = 0
-    wait = 0
 
-    for i in range(iterations):
+    for _ in range(iterations):
         indices = np.random.permutation(m)
         X_shuffled = X[indices]
         y_shuffled = y[indices]
@@ -49,10 +52,7 @@ def train_logreg(X, y, learning_rate=0.05, iterations=1500, batch_size=64):
         for j in range(0, m, batch_size):
             X_batch = X_shuffled[j:j + batch_size]
             y_batch = y_shuffled[j:j + batch_size]
-            z = np.dot(X_batch, weights)
-            h = sigmoid(z)
-            gradient = np.dot(X_batch.T, (h - y_batch)) / len(y_batch)
-
+            gradient = compute_gradient(X_batch, y_batch, weights)
             weights -= learning_rate * gradient
 
         # 计算训练集准确率
@@ -60,21 +60,16 @@ def train_logreg(X, y, learning_rate=0.05, iterations=1500, batch_size=64):
         if train_accuracy > best_accuracy:
             best_accuracy = train_accuracy
             best_weights = weights.copy()
-            wait = 0
-        else:
-            wait += 1
 
     return best_weights
 
-
 def predict_logreg(X, weights):
     z = np.dot(X, weights)
     h = sigmoid(z)
-    return np.where(h > 0.5, 1, 0)
+    return np.where(h >= 0.5, 1, 0)
 
 def load_data(filepath):
-    data = []
-    labels = []
+    data, labels = [], []
     with open(filepath, 'r') as f:
         for line in f:
             line = line.strip().split()
@@ -85,17 +80,14 @@ def load_data(filepath):
 
     data = np.array(data)
     labels = np.array(labels)
-
     target_variable = np.where((labels == 2) | (labels == 3), 0, labels)
 
     return data, target_variable
 
-
 def accuracy(X, y, weights):
     predictions = predict_logreg(X, weights)
     return np.mean(predictions == y)
 
-
 if __name__ == "__main__":
     X_train, y_train = load_data("./horse-colic.data")
     X_test, y_test = load_data("./horse-colic.test")
@@ -109,3 +101,4 @@ if __name__ == "__main__":
 
     final_accuracy = accuracy(X_test, y_test, final_weights)
     print(f"最终测试集的准确率: {final_accuracy:.4f}")
+