Update engram.py

engram.py

@@ -75,25 +75,4 @@ class EngramLayer(nn.Module):
         conv_out = conv_out.permute(0, 2, 1)
     
         output = F.silu(conv_out) + gated_value
-        return output
-
-
-if __name__=="__main__":
-    from transformers import AutoTokenizer
-
-    from compression import VocabCompressor
-    from config import EngramConfig
-
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
-    config = EngramConfig()
-    compressor = VocabCompressor(tokenizer)
-    mapping = compressor.build_mapping()
-    engram_layer = EngramLayer(config, mapping)
-
-    input_ids = torch.randint(0, 1000, (1, 10))
-    hidden_state = torch.randn(1, 10, 768)
-
-    out = engram_layer(input_ids, hidden_state)
-
-    print(engram_layer)
-    print("Output shape: ", out.shape)
+        return output