formatting

custom_generate/generate.py

@@ -11,7 +11,10 @@ from transformers.generation.logits_process import (
     TopKLogitsWarper,
     TopPLogitsWarper,
 )
-from transformers.generation.utils import GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput
+from transformers.generation.utils import (
+    GenerateDecoderOnlyOutput,
+    GenerateEncoderDecoderOutput,
+)
 
 
 def generate(
@@ -50,8 +53,12 @@ def generate(
     # 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)
-    threshold = getattr(generation_config, "threshold", 17.0)  # Default threshold for confidence (positive value)
-    conf_topk = getattr(generation_config, "conf_topk", 20)  # Number of top tokens for confidence calculation
+    threshold = getattr(
+        generation_config, "threshold", 17.0
+    )  # Default threshold for confidence (positive value)
+    conf_topk = getattr(
+        generation_config, "conf_topk", 20
+    )  # Number of top tokens for confidence calculation
 
     # If DeepCONF is not enabled, fall back to standard sampling
     if not enable_conf:
@@ -83,16 +90,26 @@ def generate(
     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_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)
+    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:
+    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()
@@ -101,7 +118,13 @@ def generate(
         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
+            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))
 
@@ -110,22 +133,36 @@ def generate(
     raw_logits = () if (return_dict_in_generate and output_logits) else None
     decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
     cross_attentions = () if (return_dict_in_generate and output_attentions) else None
-    decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+    decoder_hidden_states = (
+        () if (return_dict_in_generate and output_hidden_states) else None
+    )
 
     # If model is an encoder-decoder, retrieve encoder attention weights and hidden states
     if return_dict_in_generate and model.config.is_encoder_decoder:
-        encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
-        encoder_hidden_states = model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+        encoder_attentions = (
+            model_kwargs["encoder_outputs"].get("attentions")
+            if output_attentions
+            else None
+        )
+        encoder_hidden_states = (
+            model_kwargs["encoder_outputs"].get("hidden_states")
+            if output_hidden_states
+            else None
+        )
 
     # Keep track of which sequences are already finished
     batch_size, cur_len = input_ids.shape[:2]
-    unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+    unfinished_sequences = torch.ones(
+        batch_size, dtype=torch.long, device=input_ids.device
+    )
     # Use public kv-cache via past_key_values
 
     # Initialize confidence tracking
     # Use deque for sliding window with fixed size
     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
+    conf_grouped_sums = [
+        0.0 for _ in range(batch_size)
+    ]  # Running sums for efficient mean calculation
 
     # Optional per-step confidences for debugging/visualization
     step_confidences = [] if (return_dict_in_generate and output_confidences) else None
@@ -141,8 +178,14 @@ def generate(
         model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)
 
         # Prepare variable output controls
-        model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
-        model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
+        model_inputs.update(
+            {"output_attentions": output_attentions} if output_attentions else {}
+        )
+        model_inputs.update(
+            {"output_hidden_states": output_hidden_states}
+            if output_hidden_states
+            else {}
+        )
 
         # Forward pass with proper KV cache handling
         with torch.no_grad():
@@ -181,14 +224,18 @@ def generate(
                 raw_logits += (next_token_logits,)
             if output_attentions:
                 decoder_attentions += (
-                    (outputs.decoder_attentions,) if model.config.is_encoder_decoder else (outputs.attentions,)
+                    (outputs.decoder_attentions,)
+                    if model.config.is_encoder_decoder
+                    else (outputs.attentions,)
                 )
                 if model.config.is_encoder_decoder:
                     cross_attentions += (outputs.cross_attentions,)
 
             if output_hidden_states:
                 decoder_hidden_states += (
-                    (outputs.decoder_hidden_states,) if model.config.is_encoder_decoder else (outputs.hidden_states,)
+                    (outputs.decoder_hidden_states,)
+                    if model.config.is_encoder_decoder
+                    else (outputs.hidden_states,)
                 )
 
         # Token selection
@@ -203,8 +250,12 @@ def generate(
         # This uses the raw logits (next_token_logits) before warpers are applied.
         probs = F.softmax(next_token_logits, dim=-1)
 
-        deepconf_stopping = torch.ones(batch_size, dtype=torch.bool, device=input_ids.device)
-        step_conf_values = [0.0] * batch_size  # collect per-sequence confidences for this step (full batch)
+        deepconf_stopping = torch.ones(
+            batch_size, dtype=torch.bool, device=input_ids.device
+        )
+        step_conf_values = [
+            0.0
+        ] * batch_size  # collect per-sequence confidences for this step (full batch)
 
         for i in range(batch_size):
             if not unfinished_sequences[i]:
@@ -233,11 +284,15 @@ def generate(
 
         if step_confidences is not None:
             # Store this step's confidences as a tensor of shape (batch,)
-            step_confidences.append(torch.tensor(step_conf_values, device=input_ids.device))
+            step_confidences.append(
+                torch.tensor(step_conf_values, device=input_ids.device)
+            )
 
         # Finished sentences should have their next token be a padding token
         if has_eos_stopping_criteria and pad_token_id is not None:
-            next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+            next_tokens = next_tokens * unfinished_sequences + pad_token_id * (
+                1 - unfinished_sequences
+            )
 
         # Update generated ids, model inputs, and length for next step
         input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
@@ -245,7 +300,11 @@ def generate(
         if model_kwargs.get("attention_mask") is not None:
             attn = model_kwargs["attention_mask"]
             model_kwargs["attention_mask"] = torch.cat(
-                [attn, torch.ones((batch_size, 1), dtype=attn.dtype, device=attn.device)], dim=-1
+                [
+                    attn,
+                    torch.ones((batch_size, 1), dtype=attn.dtype, device=attn.device),
+                ],
+                dim=-1,
             )
         # Update cache_position for next step (single next token)
         model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + 1