Hey everyone, I have a code in python that creates and train a TabTransformer model that I, then, extract its embeddings so I can use it to generate embedding vectors for tabular data that I have.
Here is the code:
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
# Define the TabTransformer model
class TabTransformer(nn.Module):
def __init__(self, num_features, num_classes, dim_embedding=64, num_heads=4, num_layers=4):
super(TabTransformer, self).__init__()
self.embedding = nn.Linear(num_features, dim_embedding)
encoder_layer = nn.TransformerEncoderLayer(d_model=dim_embedding, nhead=num_heads, batch_first=True)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.classifier = nn.Linear(dim_embedding, num_classes)
def forward(self, x):
x = self.embedding(x)
x = x.unsqueeze(1) # Adding a sequence length dimension
x = self.transformer(x)
x = torch.mean(x, dim=1) # Pooling
x = self.classifier(x)
return x
data = pd.DataFrame({
"customer_id": [1, 2, 3, 4, 5],
"product_category": [1, 2, 1, 3, 2],
"ammount": [100, 200, 150, 300, 250]
})
X = data.drop("customer_id", axis = 1)
y = data["customer_id"] - 1
# Splitting the dataset into training and test sets
# X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.0, random_state=42)
X_train = X
X_test = X
y_train = y
y_test = y
# Standardizing the features
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
# Model parameters
num_features = X_train_scaled.shape[1]
num_classes = 5 # Adjusted based on unique customer ids
# Initialize the model, loss, and optimizer
model = TabTransformer(num_features, num_classes).to(torch.device('cpu'))
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# Converting data to tensors
X_train_tensor = torch.FloatTensor(X_train_scaled)
y_train_tensor = torch.LongTensor(y_train.values)
# Training loop
for epoch in range(100):
optimizer.zero_grad()
output = model(X_train_tensor)
loss = criterion(output, y_train_tensor)
loss.backward()
optimizer.step()
if epoch % 10 == 0:
print(f'Epoch {epoch}, Loss: {loss.item()}')
# Evaluation
model.eval()
X_test_tensor = torch.FloatTensor(X_test_scaled)
y_test_tensor = torch.LongTensor(y_test.values)
with torch.no_grad():
predictions = model(X_test_tensor)
_, predicted_classes = torch.max(predictions, 1)
accuracy = (predicted_classes == y_test_tensor).float().mean()
print(f'Test Accuracy: {accuracy.item()}')
# Get embedding from test data
embedding = model.embedding
embedding(X_test_tensor)
This works great, but now I want to convert this code to Elixir so I can use it in my backend to serve embeddings without needing to connect to some python code/service.
I’m very new to the ML world, so I’m sincerely not sure how I should start with doing that, especially the part that actually creates the model itself.
Update
So, I started trying to convert it, so far I was able to generate the inputs and the scaler, but I’m totally stuck in the model creation, I have no idea what are the equivalent functions in Axon for the ones used in the TabTransformer class.
Mix.install([
{:kino_explorer, "~> 0.1.20"},
{:axon, "~> 0.7"},
{:scholar, github: "elixir-nx/scholar"},
{:table_rex, "~> 4.0", override: true}
])
alias Explorer.{Series, DataFrame}
alias Scholar.Preprocessing.StandardScaler
require DataFrame
require Series
data = DataFrame.new(
customer_id: [1, 2, 3, 4, 5],
product_category: [1, 2, 1, 3, 2],
amount: [100, 200, 150, 300, 250]
)
x = DataFrame.discard(data, :customer_id)
y =
data
|> DataFrame.select(:customer_id)
|> DataFrame.mutate(customer_id: customer_id - 1)
x_train = Nx.stack(x, axis: 1)
y_train = Nx.stack(y, axis: 1)
x_test = Nx.stack(x, axis: 1)
y_test = Nx.stack(y, axis: 1)
scaler = StandardScaler.fit(x_train, axes: [0])
x_train_scaled = StandardScaler.transform(scaler, x_train)
x_test_scaled = StandardScaler.transform(scaler, x_test)
{_, num_features} = Nx.shape(x_train_scaled)
# Adjusted based on unique customer ids
num_classes = 5
