synthetic-np / SyntheticModel.py
fivedollarwitch's picture
Upload 5 files
26ef0d6 verified
import torch
import torch.nn as nn
from torch.utils.data import Dataset
# PyTorch Dataset
class RealEstateDataset(Dataset):
def __init__(self, df, categorical_columns, numerical_columns):
self.X_cat = torch.tensor(df[categorical_columns].values, dtype=torch.long)
self.X_num = torch.tensor(df[numerical_columns].values, dtype=torch.float32)
self.y_price = torch.tensor(df['sale_price'].values, dtype=torch.float32)
self.y_date = torch.tensor(df['sale_date'].values, dtype=torch.long)
def __len__(self):
return len(self.X_num)
def __getitem__(self, idx):
return self.X_cat[idx], self.X_num[idx], self.y_price[idx], self.y_date[idx]
# Model Definition
class RealEstateModel(nn.Module):
def __init__(self, num_numerical, embedding_dims, num_date_classes):
super(RealEstateModel, self).__init__()
# Embeddings for categorical variables
self.embeddings = nn.ModuleList([
nn.Embedding(num_classes, dim)
for num_classes, dim in embedding_dims
])
# Fully connected layers
embedding_dim_sum = sum(dim for _, dim in embedding_dims)
self.fc1 = nn.Linear(embedding_dim_sum + num_numerical, 256)
self.fc2 = nn.Linear(256, 128)
# Output layers
self.price_head = nn.Linear(128, 1)
self.date_head = nn.Linear(128, num_date_classes)
def forward(self, x_cat, x_num):
# Embedding layers
x = [emb(x_cat[:, i]) for i, emb in enumerate(self.embeddings)]
x = torch.cat(x, dim=1)
# Concatenate with numerical features
x = torch.cat([x, x_num], dim=1)
# Fully connected layers
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
# Output heads
price = self.price_head(x)
sale_date = self.date_head(x)
return price, sale_date