Idea: move schema into .inputs.base._Inputs class

torecsys/inputs/base/__init__.py

@@ -1,8 +1,10 @@
 r"""torecsys.inputs.base is a sub-module of base inputs class.
 """
 
+from collections import namedtuple
 import torch.nn as nn
 
+
 class _Inputs(nn.Module):
     r"""General Input class.
     """
@@ -17,6 +19,12 @@ def __len__(self) -> int:
             int: Size of embedding tensor, or Number of inputs' fields.
         """
         return self.length
+
+    def get_schema(self) -> namedtuple:
+        return self.schema
+
+    def set_schema(self, *args):
+        raise NotImplementedError("set_schema cannot be called in the base class.")
 
 
 # from .audio_inp import AudioInputs

torecsys/inputs/base/value_inp.py

@@ -1,6 +1,8 @@
 from . import _Inputs
-from torecsys.utils.decorator import jit_experimental
+from collections import namedtuple
 import torch
+from torecsys.utils.decorator import jit_experimental, no_jit_experimental_by_namedtensor
+from typing import List
 
 class ValueInputs(_Inputs):
     r"""Base Inputs class for value to be passed directly.
@@ -10,8 +12,8 @@ class ValueInputs(_Inputs):
     #. add transforms for value inputs to do preprocessing
 
     """
-    @jit_experimental
-    def __init__(self, num_fields: int):
+    @no_jit_experimental_by_namedtensor
+    def __init__(self, inp_fields: List[str]):
         r"""Initialize ValueInputs
         
         Args:
@@ -23,8 +25,15 @@ def __init__(self, num_fields: int):
         # refer to parent class
         super(ValueInputs, self).__init__()
 
-        # bind length to num_fields 
-        self.length = num_fields
+        # initialize schema
+        self.set_schema(inp_fields)
+
+        # bind length to length of inp_fields 
+        self.length = len(inp_fields)
+
+    def set_schema(self, inputs: List[str]):
+        schema = namedtuple("Schema", ["inputs"])
+        self.schema = schema(inputs=inputs)
 
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         r"""Forward calculation of ValueInputs.
@@ -33,7 +42,7 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
             inputs (T), shape = (B, N): Tensor of values in input fields.
         
         Returns:
-            T, shape = (B, 1, N): Outputs of ValueInputs
+            T, shape = (B, N): Outputs of ValueInputs
         """
-        # unsqueeze(1) and return
-        return inputs.unsqueeze(1)
+        inputs.names = ("B", "E")
+        return inputs