Untitled

def train_one_epoch(
    model: Type[LogisticRegressionModel],
    train_set: np.ndarray, 
    train_labels: np.ndarray, 
    val_set: np.ndarray, 
    val_labels:np.ndarray,
    batch_size: int, 
    learning_rate: float, 
    validation_every_x_step: int
) -> float:
    """
    Trains the model for one epoch on the entire dataset with the given learning rate and batch size
    
    Args:
        model (class): the model used to train
        train_set (np.ndarray): Numpy array of shape [val_size x num_features]
        train_labels (np.ndarray): Numpy array of shape [val_size]
        val_set (np.ndarray): Numpy array of shape [val_size x num_features]
        val_labels (np.ndarray): Numpy array of shape [val_size]
        batch_size (int): the batch size to be used to iterate through the dataset
        learning_rate (float): the learning rate to be used for mini-batch gradient descent optimization
        validation_every_x_step (int): the number of steps to wait before performing validation
    
    Returns:
        train_losses (list): a list of training losses
        train_accuracies (list): a list of training accuracies
        # train_steps (list): a list of the training batch ids, i.e. each element is the n-th batch of the training set
        train_steps (list): a list of the number of training steps. One training step is defined as one forward pass 
                            (i.e. calculating the loss) AND one backward pass (i.e. calculating the gradients and updating the parameters) 
                            of a mini-batch of samples through the model
        val_losses (list): a list of validation losses
        val_accuracies (list): a list of validation accuracies
        val_steps (list): a list of the validation steps. One validation step is defined one forward pass of the validaton mini-batch
                            samples through the model
    """
    train_losses = []
    train_accuracies = []
    train_steps = []
    val_losses = []
    val_accuracies = []
    val_steps = []
    step_count = 0

    # Iterate through the training set and append the corresponding metrics to the list
    for x_batch, targets in iterate_samples(batch_size, train_set, train_labels, True):
        
        step_count += 1
        pred = model.__call__(x_batch)
        grad_W, grad_b = compute_gradients(x_batch, pred, targets)

        model.W = model.W - learning_rate * grad_W
        model.b = model.b - learning_rate * grad_b

        train_loss, train_acc = validation(model, train_set, train_labels, batch_size)
        train_losses.append(train_loss)
        train_accuracies.append(train_acc)
        train_steps.append(step_count)

        # perform validation depending on the value of `validation_every_x_step`
        if (step_count % validation_every_x_step) == 0 or step_count == 1:
            
            validation_loss, validation_acc = validation(model, val_set, val_labels, batch_size)
            val_losses.append(validation_loss)
            val_accuracies.append(validation_acc)
            val_steps.append(step_count)
        
    return train_losses, train_accuracies, train_steps, val_losses, val_accuracies, val_steps