kjbk

modelLM.py

@@ -29,30 +29,30 @@ class OBILanguageModel(PreTrainedModel):
 
        
     
-        def forward(self, idx, targets=None):
-            tok_emb = self.token_embedding_table(idx)
-            pos_emb = None  # Initialize pos_emb to None
-            try:
-                pos_emb = self.position_embedding_table(torch.arange(idx.size(1), device='cpu'))
-            except IndexError as e:
-                # Handle the IndexError by initializing pos_emb with zeros
-                print(f"IndexError: {e}")
-                print(f"idx.size(1): {idx.size(1)}")
-                print(f"Positional embedding table shape: {self.position_embedding_table.weight.shape}")
-                pos_emb = torch.zeros((idx.size(1), self.config.hidden_size), device=device)
+    def forward(self, idx, targets=None):
+        tok_emb = self.token_embedding_table(idx)
+        pos_emb = None  # Initialize pos_emb to None
+        try:
+            pos_emb = self.position_embedding_table(torch.arange(idx.size(1), device='cpu'))
+        except IndexError as e:
+            # Handle the IndexError by initializing pos_emb with zeros
+            print(f"IndexError: {e}")
+            print(f"idx.size(1): {idx.size(1)}")
+            print(f"Positional embedding table shape: {self.position_embedding_table.weight.shape}")
+            pos_emb = torch.zeros((idx.size(1), self.config.hidden_size), device=device)
 
-            x = tok_emb + pos_emb
-            x = self.transformer(x, x)
-            x = self.ln1(x)
-            x = self.ln2(x)
-            logits = self.lm_head(x)
+        x = tok_emb + pos_emb
+        x = self.transformer(x, x)
+        x = self.ln1(x)
+        x = self.ln2(x)
+        logits = self.lm_head(x)
 
-            if targets is None:
-                loss = None
-            else:
-                loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1))
+        if targets is None:
+            loss = None
+        else:
+            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1))
 
-            return logits, loss
+        return logits, loss
 
 
     def generate(self, idx, max_new_tokens):