from sklearn.datasets import fetch_california_housing

housing = fetch_california_housing(as_frame=True)

from sklearn.model_selection import train_test_split

X_train_full, X_test, y_train_full, y_test = train_test_split(housing['data'], housing['target'], test_size=0.25, random_state=42)
X_train, X_valid, y_train, y_valid = train_test_split(X_train_full, y_train_full, test_size=0.25, random_state=42)

X_means, X_stds = X_train.mean(axis=0), X_train.std(axis=0)
X_train = (X_train - X_means) / X_stds
X_valid = (X_valid - X_means) / X_stds
X_test = (X_test - X_means) / X_stds

import torch

device = torch.device("cpu")

import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin

class MLP(nn.Module, PyTorchModelHubMixin):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(8, 50)
        self.fc2 = nn.Linear(50, 50)
        self.fc3 = nn.Linear(50, 50)
        self.fc4 = nn.Linear(50, 1)

    def forward(self, x):
        act = torch.relu(self.fc1(x))
        act = torch.relu(self.fc2(act))
        act = torch.relu(self.fc3(act))
        return self.fc4(act)

model = MLP.from_pretrained("sadhaklal/mlp-california-housing")
model.to(device)
model.eval()

# Let's predict on 3 unseen examples from the test set:
print(f"Ground truth housing prices: {y_test.values[:3]}")
x_new = torch.tensor(X_test.values[:3], dtype=torch.float32)
x_new = x_new.to(device)
with torch.no_grad():
    preds = model(x_new)
print(f"Predicted housing prices: {preds.squeeze()}")