Adapt generate for text-generation pipeline

modeling_logos.py

@@ -2,12 +2,13 @@
 
 from __future__ import annotations
 
+from copy import deepcopy
 import importlib
 from typing import Any, Dict, Optional
 
 import torch
 from torch import nn
-from transformers import PreTrainedModel
+from transformers import GenerationConfig, PreTrainedModel
 try:
     from transformers.generation import GenerationMixin
 except Exception:  # pragma: no cover - older/newer transformers layout guard
@@ -31,6 +32,18 @@ except ImportError:
     LogosTransformer = importlib.import_module("models.logos").LogosTransformer
 
 
+_GENERATION_CONFIG_KEYS = set(GenerationConfig().to_dict())
+_SAMPLING_ONLY_KEYS = (
+    "temperature",
+    "top_k",
+    "top_p",
+    "min_p",
+    "typical_p",
+    "epsilon_cutoff",
+    "eta_cutoff",
+)
+
+
 def _reorder_cache_value(value: Any, beam_idx: torch.LongTensor) -> Any:
     if isinstance(value, torch.Tensor):
         if value.dim() > 0 and value.size(0) == beam_idx.size(0):
@@ -95,6 +108,24 @@ class LogosForCausalLM(PreTrainedModel, GenerationMixin):
         self.model.lm_head.weight = self.model.token_emb.weight
         self._reset_rotary_buffers()
 
+    def generate(self, *args: Any, **kwargs: Any):
+        generation_config = kwargs.get("generation_config")
+        if generation_config is not None:
+            explicit_keys = _GENERATION_CONFIG_KEYS.intersection(kwargs)
+            if explicit_keys:
+                merged = deepcopy(generation_config)
+                for key in sorted(explicit_keys):
+                    setattr(merged, key, kwargs.pop(key))
+                if "max_new_tokens" in explicit_keys and "max_length" not in explicit_keys:
+                    merged.max_length = None
+                if getattr(merged, "do_sample", None) is False:
+                    for key in _SAMPLING_ONLY_KEYS:
+                        if key not in explicit_keys and hasattr(merged, key):
+                            setattr(merged, key, None)
+                kwargs["generation_config"] = merged
+
+        return super().generate(*args, **kwargs)
+
     def _reset_rotary_buffers(self) -> None:
         for module in self.modules():
             if module.__class__.__name__ != "RotaryEmbedding":