Project: 023_Rethinking_Reward_Modeling_in_Preference-based_Lar | Function: train_model (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements the unified training procedure for both Bradley-Terry and classification-based reward models with order consistency objective, supporting both 'clf' and 'siamese' modes for different reward modeling approaches

**KEY TERMS:** {
  "model_mode": "reward model type (BT or CLF)",
  "forward_mlp": "classification approach",
  "forward_siamese": "Bradley-Terry approach",
  "BCELoss": "binary cross-entropy for order consistency"
}

**CONTEXT CODE:** 
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split

class MLP(nn.Module):
    def __init__(self, input_dim):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(input_dim, 1024)  # First hidden layer
        self.fc2 = nn.Linear(1024, 512)         # Second hidden layer
        self.fc3 = nn.Linear(512, 1)           # Output layer (binary classification)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)  # Sigmoid activation for binary classification
        return x

def forward_mlp(model, input):
    return torch.sigmoid(model(input))

def forward_siamese(model, input1, input2):
    raise NotImplementedError("Function implementation has been masked for testing")
def train_model(model, device, model_mode, X_train, y_train, X_test, y_test, epochs=3, lr=0.001, batch_size = 512, verbose=False, ealry_stopping_patience=3):
    raise NotImplementedError("Function implementation has been masked for testing")
def save_model(model, path):
    torch.save(model.state_dict(), path)


Complete the function implementation replacing "raise NotImplementedError".