Update modeling_pragya.py

modeling_pragya.py

@@ -1,6 +1,7 @@
 import dataclasses
 
 import torch
+import torch.nn.functional as F
 from transformers import GenerationMixin, PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
@@ -86,7 +87,8 @@ class PragyaForCausalLM(PreTrainedModel, GenerationMixin):
     def get_output_embeddings(self):
         return self.model.lm_head
 
-    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+    def forward(self, input_ids=None, attention_mask=None, labels=None,
+                logits_to_keep=0, **kwargs):
         if attention_mask is not None and not bool(torch.all(attention_mask == 1)):
             # The model has no attention-mask support. Padded batches would silently
             # attend to pad tokens, so refuse rather than return wrong results.
@@ -94,10 +96,27 @@ class PragyaForCausalLM(PreTrainedModel, GenerationMixin):
                 "attention_mask with padding is not supported. Run one sequence at a time "
                 "or pass unpadded batches of equal length."
             )
-        # With labels: full logits (B, T, V) and loss.
-        # Without labels: logits for the last position only, (B, 1, V).
-        logits, loss = self.model(input_ids, targets=labels)
-        return CausalLMOutputWithPast(loss=loss, logits=logits)
+        if labels is not None:
+            # Training path: GPT computes the loss itself.
+            logits, loss = self.model(input_ids, targets=labels)
+            return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+        # Inference path. logits_to_keep=0 means logits for every position (B, T, V),
+        # which lm-evaluation-harness and loglikelihood scoring need. logits_to_keep=k
+        # returns the last k positions. Hidden states are captured after ln_f, and the
+        # projection is the same F.linear(x, lm_head.weight) GPT uses internally.
+        captured = {}
+        handle = self.model.transformer.ln_f.register_forward_hook(
+            lambda module, inp, out: captured.__setitem__("x", out))
+        try:
+            self.model(input_ids)
+        finally:
+            handle.remove()
+        x = captured["x"]
+        if logits_to_keep:
+            x = x[:, -logits_to_keep:, :]
+        logits = F.linear(x, self.model.lm_head.weight)
+        return CausalLMOutputWithPast(logits=logits)
 
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
         # No KV cache: GPT keeps its own cache format, so every step runs on the full sequence.