Implement GPT2 class

Author: satojkovic

pytorch/gpt_pytorch.py

@@ -3,6 +3,8 @@
 import sys
 import os
 from utils import load_encoder_hparams_and_params
+from model import GPT2
+import torchsummaryX
 
 
 if __name__ == "__main__":
@@ -26,3 +28,6 @@
     input_ids = encoder.encode(args.prompt)
     input_text = encoder.decode(input_ids)
     print("input_ids:", input_ids)
+
+    model = GPT2(params, hparams, drop_p=0.1)
+    torchsummaryX.summary(model, torch.ones(1, len(input_ids), dtype=torch.long))

pytorch/model.py

@@ -60,11 +60,32 @@ def forward(self, x):
 
 
 class GPT2(nn.Module):
-    def __init__(self):
+    def __init__(self, params, hparams, drop_p=0.1):
         super().__init__()
-
-    def forward(self, x):
-        pass
+        self.params = params
+        self.hparams = hparams
+        self.drop_p = drop_p
+        self.h_dim = hparams["n_embd"]
+        self.n_heads = hparams["n_head"]
+
+        self.wte = self.params["wte.weight"]
+        self.wpe = self.params["wpe.weight"]
+
+        self.blocks = []
+        for _ in range(self.hparams["n_layer"]):
+            block = TransformerDecoderBlock(
+                h_dim=self.h_dim, n_heads=self.n_heads, drop_p=self.drop_p
+            )
+            self.blocks.append(block)
+        self.layer_norm = nn.LayerNorm(self.h_dim)
+
+    def forward(self, input_ids):
+        x = self.wte[input_ids] + self.wpe[list(range(input_ids.shape[1]))]
+        for block in self.blocks:
+            x = block(x)
+        x = self.layer_norm(x)
+        out = x @ self.wte.T
+        return out
 
 
 if __name__ == "__main__":

pytorch/utils.py

@@ -8,7 +8,7 @@
 # import picoGPT
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
 sys.path.append(os.path.join(os.path.dirname(__file__), "../picoGPT"))
-from encoder import get_encoder
+from picoGPT.encoder import get_encoder
 
 
 def save_file(file_path, req):