From 748ee9626fa9b8b97aa32c4fa31efd1718bd69ed Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Fri, 1 Nov 2019 15:33:53 +0100
Subject: [PATCH] added FeatureAssembler

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

diff --git a/pts/modules/feature.py b/pts/modules/feature.py
index a13351f..dae3244 100644
--- a/pts/modules/feature.py
+++ b/pts/modules/feature.py
@@ -36,5 +36,54 @@ class FeatureEmbedder(nn.Module):
             in zip(self.__embedders, cat_feature_slices)
         ], dim=-1)
 
+
 class FeatureAssembler(nn.Module):
-    pass
\ No newline at end of file
+    def __init__(T: int,
+                 use_static_cat: bool = False,
+                 use_static_real: bool = False,
+                 use_dynamic_cat: bool = False,
+                 use_dynamic_real: bool = False,
+                 embed_static: Optional[FeatureEmbedder] = None,
+                 embed_dynamic: Optional[FeatureEmbedder] = None,
+                 dtype: torch.dtype = torch.float32) -> None:
+        super().__init__()
+
+        self.T = T
+        self.use_static_cat = use_static_cat
+        self.use_static_real = use_static_real
+        self.use_dynamic_cat = use_dynamic_cat
+        self.use_dynamic_real = use_dynamic_real
+
+        self.embed_static: Callable[[torch.Tensor], torch.
+                                    Tensor] = embed_static or (lambda x: x)
+        self.embed_dynamic: Callable[[torch.Tensor], torch.
+                                     Tensor] = embed_dynamic or (lambda x: x)
+
+    def forward(
+            self,
+            feat_static_cat: torch.Tensor,
+            feat_static_real: torch.Tensor,
+            feat_dynamic_cat: torch.Tensor,
+            feat_dynamic_real: torch.Tensor,
+    ) -> torch.Tensor:
+        processed_features = [
+            self.process_static_cat(feat_static_cat),
+            self.process_static_real(feat_static_real),
+            self.process_dynamic_cat(feat_dynamic_cat),
+            self.process_dynamic_real(feat_dynamic_real),
+        ]
+
+        return torch.cat(processed_features, dim=-1)
+
+    def process_static_cat(self, feature: torch.Tensor) -> torch.Tensor:
+        feature = self.embed_static(feature.to(self.dtype))
+        return feature.unsqueeze(1).expand(-1, self.T, -1)
+
+    def process_dynamic_cat(self, feature: torch.Tensor) -> torch.Tensor:
+        return self.embed_dynamic(feature.to(self.dtype))
+
+    def process_static_real(self, feature: torch.Tensor) -> torch.Tensor:
+        return feature.unsqueeze(1).expand(-1, self.T, -1)
+
+    def process_dynamic_real(self, feature: torch.Tensor) -> torch.Tensor:
+        return feature