Update custom_generate/generate.py

custom_generate/generate.py

@@ -1,6 +1,7 @@
 from collections import deque
 from typing import Any, Optional, Union
 
+import numpy as np
 import torch
 import torch.nn.functional as F
 
@@ -40,6 +41,12 @@ def generate(
         depending on `return_dict_in_generate` and model type.
     """
 
+    # Ensure processors/criteria are defined
+    if logits_processor is None:
+        logits_processor = LogitsProcessorList()
+    if stopping_criteria is None:
+        stopping_criteria = StoppingCriteriaList()
+
     # Get DeepCONF parameters from generation_config or set defaults
     enable_conf = getattr(generation_config, "enable_conf", False)
     window_size = getattr(generation_config, "window_size", 2048)
@@ -74,9 +81,29 @@ def generate(
     output_logits = generation_config.output_logits
     return_dict_in_generate = generation_config.return_dict_in_generate
     output_confidences = getattr(generation_config, "output_confidences", False)
+    # Optional DeepConf variant helpers (compute threshold from warmup confidences)
+    deepconf_variant = getattr(generation_config, "deepconf_variant", None)  # "low" or "high"
+    deepconf_eta = getattr(generation_config, "deepconf_eta", None)  # float in (0,1)
+    deepconf_warmup_confidences = getattr(generation_config, "deepconf_warmup_confidences", None)  # list/1D tensor
     has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria)
     do_sample = generation_config.do_sample
 
+    # If a variant is requested and a warmup set of confidences is provided, derive the threshold
+    if enable_conf and threshold is not None:
+        pass
+    elif enable_conf and deepconf_variant is not None and deepconf_warmup_confidences is not None:
+        confs = deepconf_warmup_confidences
+        if hasattr(confs, "detach"):
+            confs = confs.detach().cpu().numpy()
+        elif isinstance(confs, torch.Tensor):
+            confs = confs.cpu().numpy()
+        confs = np.asarray(confs, dtype=np.float32).ravel()
+        eta = deepconf_eta
+        if eta is None:
+            eta = 0.1 if deepconf_variant == "low" else 0.9 if deepconf_variant == "high" else 0.5
+        pct = max(0.0, min(100.0, 100.0 - (eta * 100.0)))
+        threshold = float(np.percentile(confs, pct))
+
     # Initialize attention / hidden states / scores tuples
     scores = () if (return_dict_in_generate and output_scores) else None
     raw_logits = () if (return_dict_in_generate and output_logits) else None
@@ -99,8 +126,6 @@ def generate(
     conf_group_lists = [deque(maxlen=window_size) for _ in range(batch_size)]
     conf_grouped_sums = [0.0 for _ in range(batch_size)]  # Running sums for efficient mean calculation
 
-    # Initialize via prepare_inputs_for_generation
-
     # Optional per-step confidences for debugging/visualization
     step_confidences = [] if (return_dict_in_generate and output_confidences) else None