off

modelLM.py

@@ -1,13 +1,18 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PreTrainedModel
+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)
+        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,
@@ -20,21 +25,40 @@ 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)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)  # Use length of SentencePiece vocab
 
-    def forward(self, input_ids, attention_mask=None, labels=None):
-        tok_emb = self.token_embedding_table(input_ids)
-        pos_emb = self.position_embedding_table(torch.arange(input_ids.size(1), device='cpu'))
+       
+    
+    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 = self.transformer(x, x)
         x = self.ln1(x)
         x = self.ln2(x)
         logits = self.lm_head(x)
 
-        # Always compute the loss if labels are provided
-        loss = None
-        if labels is not None:
-            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1))
+            # Always compute the loss, and set it to None if targets are not provided
+        loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1)) if targets is not None else None
+
+        return (logits, loss)
+
 
-        return {"logits": logits, "loss": loss}
+    def generate(self, idx, max_new_tokens):
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.config.block_size:]
+            logits, loss = self(idx_cond)
+            logits = logits[:, -1, :]
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx