Update ChemQ3MTP/modeling_chemq3mtp.py

ChemQ3MTP/modeling_chemq3mtp.py

@@ -5,6 +5,7 @@
 # ========================
 
 import os
+import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -24,6 +25,7 @@ import json
 import numpy as np
 from collections import Counter
 from rdkit.Chem import rdMolDescriptors
+from transformers.generation.utils import GenerationMixin
 
 logger = logging.get_logger(__name__)
 
@@ -373,16 +375,34 @@ class ChemQ3MTPForCausalLM(Qwen2ForCausalLM):
             **kwargs
         )
 
+    def generate(self, *args, **kwargs):
+        """
+        Wrap HF GenerationMixin.generate so that
+        max_new_tokens = ceil(0.25 * prompt_length) when the caller
+        omits both max_new_tokens and max_length.
+        """
+        # only touch if user did NOT set any length cap
+        if (kwargs.get("max_new_tokens") is None
+                and kwargs.get("max_length") is None):
+
+            # locate input_ids (works when passed positionally or by keyword)
+            input_ids = kwargs.get("input_ids", args[0] if args else None)
+            if input_ids is not None:
+                kwargs["max_new_tokens"] = max(1, math.ceil(input_ids.shape[1] * 0.25))
+
+        # delegate to HF implementation
+        return super().generate(*args, **kwargs)
+
     def generate_with_logprobs(
         self,
         input_ids: torch.LongTensor,
-        max_new_tokens: int = 50,
+        max_new_tokens: int = None,          # ← changed default to None
         temperature: float = 1.0,
         top_k: Optional[int] = None,
         top_p: Optional[float] = None,
         do_sample: bool = True,
         return_probs: bool = True,
-        tokenizer=None,
+        tokenizer=None
     ) -> Tuple[List[str], torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
         """
         Generate sequences with log probabilities for RL training.
@@ -392,10 +412,15 @@ class ChemQ3MTPForCausalLM(Qwen2ForCausalLM):
         1. Use log_softmax instead of log(softmax) to avoid log(0) issues
         2. Correct the gather operation for non-sampling case
         3. Handle the case where filtered logits become -inf properly
+        NEW: if max_new_tokens is not given it is set to ceil(0.25 * prompt_length).
         """
         self.eval()
         device = input_ids.device
 
+        # ---------- auto-compute max_new_tokens ----------
+        if max_new_tokens is None:
+            max_new_tokens = max(1, math.ceil(input_ids.size(1) * 0))
+        
         # Normalize input shapes
         if input_ids.dim() == 1:
             input_ids = input_ids.unsqueeze(0)