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