From 60b4ea3cba3cdc7bbe7975dc467f8f98f8b4a658 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Fri, 1 Nov 2019 14:48:55 +0100
Subject: [PATCH] FeatureEmbedder

---
 pts/modules/feature.py | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/pts/modules/feature.py b/pts/modules/feature.py
index d95cefe..a13351f 100644
--- a/pts/modules/feature.py
+++ b/pts/modules/feature.py
@@ -1,2 +1,40 @@
+import torch
+import torch.nn as nn
+
+
 class FeatureEmbedder(nn.Module):
+    def __init__(
+            self,
+            cardinalities: List[int],
+            embedding_dims: List[int],
+    ) -> None:
+        super().__init__()
+
+        self.__num_features = len(cardinalities)
+
+        def create_embedding(c: int, d: int) -> nn.Embedding:
+            embedding = nn.Embedding(c, d)
+            return embedding
+
+        self.__embedders = nn.ModuleList([
+            create_embedding(c, d)
+            for c, d in zip(cardinalities, embedding_dims)
+        ])
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        if self.__num_features > 1:
+            # we slice the last dimension, giving an array of length
+            # self.__num_features with shape (N,T) or (N)
+            cat_feature_slices = torch.chunk(features,
+                                             self.__num_features,
+                                             dim=-1)
+        else:
+            cat_feature_slices = [features]
+
+        return torch.cat([
+            embed(cat_feature_slice.squeeze(-1)) for embed, cat_feature_slice
+            in zip(self.__embedders, cat_feature_slices)
+        ], dim=-1)
+
+class FeatureAssembler(nn.Module):
     pass
\ No newline at end of file