hgvj

modelLM.py

@@ -12,7 +12,7 @@ class OBILanguageModel(PreTrainedModel):
         super(OBILanguageModel,self).__init__(config)
         self.token_embedding_table = nn.Embedding(config.vocab_size, config.hidden_size)  # Use length of SentencePiece vocab
         self.position_embedding_table = nn.Embedding(config.block_size, config.hidden_size)
-    
+  
 
         self.transformer = nn.Transformer(
             d_model=config.hidden_size,
@@ -29,29 +29,31 @@ class OBILanguageModel(PreTrainedModel):
 
        
     
-    def forward(self, idx, targets=None):
-        tok_emb = self.token_embedding_table(idx)
-        # pos_emb = self.position_embedding_table(torch.arange(idx.size(1), device=device))
-        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:
-            # Print relevant information for debugging
-            print(f"IndexError: {e}")
-            print(f"idx.size(1): {idx.size(1)}")
-            print(f"Positional embedding table shape: {self.position_embedding_table.weight.shape}")
-        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))
-        
-        return logits, loss
+        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)
+
+            if targets is None:
+                loss = None
+            else:
+                loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1))
+
+            return logits, loss
+
 
     def generate(self, idx, max_new_tokens):
         for _ in range(max_new_tokens):