Remove extra info token pathway

backbones.py

@@ -247,57 +247,6 @@ class GPTBlock(nn.Module):
         return x
 
 
-class TokenAutoDiscretization(nn.Module):
-    def __init__(
-        self,
-        n_cont_types: int,
-        n_bins: int,
-        n_embd: int,
-    ):
-        super().__init__()
-        if n_cont_types <= 0:
-            raise ValueError(f"n_cont_types must be > 0, got {n_cont_types}")
-        if n_bins <= 1:
-            raise ValueError(f"n_bins must be > 1, got {n_bins}")
-        if n_embd <= 0:
-            raise ValueError(f"n_embd must be > 0, got {n_embd}")
-
-        self.n_cont_types = n_cont_types
-        self.n_bins = n_bins
-        self.n_embd = n_embd
-        self.weight = nn.Parameter(torch.empty(n_cont_types, n_bins))
-        self.bias = nn.Parameter(torch.empty(n_cont_types, n_bins))
-        self.bin_emb = nn.Parameter(torch.empty(n_cont_types, n_bins, n_embd))
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        nn.init.normal_(self.weight, mean=0.0, std=0.02)
-        nn.init.zeros_(self.bias)
-        nn.init.normal_(self.bin_emb, mean=0.0, std=0.02)
-
-    def forward(
-        self,
-        cont_type_idx: torch.LongTensor,  # (N,)
-        value: torch.Tensor,              # (N,)
-    ) -> torch.Tensor:
-        if cont_type_idx.dim() != 1:
-            raise ValueError(
-                f"cont_type_idx must be 1D, got {tuple(cont_type_idx.shape)}"
-            )
-        if value.dim() != 1:
-            raise ValueError(f"value must be 1D, got {tuple(value.shape)}")
-        if cont_type_idx.numel() != value.numel():
-            raise ValueError("cont_type_idx and value must have the same length")
-
-        w = self.weight[cont_type_idx]      # (N, n_bins)
-        b = self.bias[cont_type_idx]        # (N, n_bins)
-        e = self.bin_emb[cont_type_idx]     # (N, n_bins, D)
-        logits = value.unsqueeze(-1) * w + b
-        probs = torch.softmax(logits, dim=-1)
-        return torch.einsum("nb,nbd->nd", probs, e)
-
-
-
 class AgeSinusoidalEncoding(nn.Module):
 
     def __init__(self, embedding_dim: int):