From a3ded691e1746c8969f9814e3c2ae92981f30bac Mon Sep 17 00:00:00 2001
From: wassname <git@wassname.org>
Date: Sun, 12 Nov 2023 17:27:17 +0800
Subject: [PATCH] misc

---
 .vscode/launch.json       |  2 ++
 research_journal.md       |  6 ++++++
 src/models/transformer.py | 13 -------------
 3 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/.vscode/launch.json b/.vscode/launch.json
index 2648c79..0705eb8 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -20,6 +20,8 @@
                 "training.world_model.start_after_epochs=1",
                 "training.actor_critic.start_after_epochs=1",
                 "training.tokenizer.steps_per_epoch=20",
+                "training.world_model.steps_per_epoch=10",
+                "training.actor_critic.steps_per_epoch=20",
             ]
         }
     ]
diff --git a/research_journal.md b/research_journal.md
index 6be63bb..243c94c 100644
--- a/research_journal.md
+++ b/research_journal.md
@@ -96,3 +96,9 @@ Debugging:
     transfrmer
     x.shape
     torch.Size([16, 340, 256])
+
+# 2023-11-12 16:58:37
+
+So I got it training, but during imagination it passes in a single token with no past steps. But the slicer seems to need at least on block? And so I get none?
+
+hmm it's because num_kept_tokens is 16 not 1. So there should be a whole block passed in ?
diff --git a/src/models/transformer.py b/src/models/transformer.py
index 431fe73..7c0c44c 100644
--- a/src/models/transformer.py
+++ b/src/models/transformer.py
@@ -109,11 +109,7 @@ class Transformer(nn.Module):
         super().__init__()
         self.config = config
         self.model = load_pretrained_model(config)
-        # self.ln_f = nn.LayerNorm(self.model.config.vocab_size, config.embed_dim)
         self.ln_f = nn.Linear(self.model.config.vocab_size, config.embed_dim)
-        # self.drop = nn.Dropout(config.embed_pdrop)
-        # self.blocks = nn.ModuleList([Block(config) for _ in range(config.num_layers)])
-        # self.ln_f = nn.LayerNorm(config.embed_dim)
 
     def generate_empty_keys_values(self, n: int, max_tokens: int) -> KeysValues:
         device = self.ln_f.weight.device  # Assumption that all submodules are on the same device
@@ -129,15 +125,6 @@ class Transformer(nn.Module):
             output_hidden_states=True,
         )
         x = outputs.logits.to(torch.float32)
-        # TODO: we output with last dim 50304 but want 2560 (embed dim)
-        x = self.ln_f(x)
-        return x
-        
-        inputs_embeds
-        x = self.drop(sequences)
-        for i, block in enumerate(self.blocks):
-            x = block(x, None if past_keys_values is None else past_keys_values[i])
-
         x = self.ln_f(x)
         return x