ghg

modelLM.py

@@ -3,12 +3,17 @@ import torch.nn as nn
 import torch.nn.functional as F
 from transformers.modeling_utils import PreTrainedModel
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
 # Define your custom language model class
 class OBILanguageModel(PreTrainedModel):
     def __init__(self, config):
         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,
             nhead=config.num_attention_heads,
@@ -21,16 +26,21 @@ class OBILanguageModel(PreTrainedModel):
         self.ln1 = nn.LayerNorm(config.hidden_size)
         self.ln2 = nn.LayerNorm(config.hidden_size)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)  # Use length of SentencePiece vocab
+
+       
     
-    def forward(self, idx, attention_mask=None, targets=None):
+    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='cpu'))
+        # 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
-
-        # Assuming you need to add attention_mask here
-        if attention_mask is not None:
-            x *= attention_mask.unsqueeze(-1)
-
         x = self.transformer(x, x)
         x = self.ln1(x)
         x = self.ln2(x)