Fix tuple index unpacking in PyTorch engine

core/engine_torch.py

@@ -1,12 +1,13 @@
 """
 PyTorch / CUDA / CPU Inference Engine for Parallel Constrained Decoding.
-Optimized for Linux containers, Hugging Face Spaces (CPU & ZeroGPU/CUDA), and cloud environments.
+Optimized for Linux containers, Hugging Face Spaces (ZeroGPU & CUDA), and cloud environments.
 """
 
 import os
 import time
 import json
 import copy
+import re
 import threading
 from typing import Dict, Any, Generator, Optional, List, Tuple
 
@@ -28,7 +29,7 @@ _gpu_lock = threading.Lock()
 # Support Hugging Face Spaces ZeroGPU if available
 try:
     import spaces
-    gpu_decorator = spaces.GPU
+    gpu_decorator = spaces.GPU(duration=60)
 except Exception:
     def gpu_decorator(fn=None, **kwargs):
         if fn is not None:
@@ -66,42 +67,6 @@ def get_torch_engine():
     return _torch_model, _torch_tokenizer, _torch_device
 
 
-def _extract_choice_logits(
-    last_logits: torch.Tensor,
-    cands_per_field: List[List[List[int]]],
-    temperature: float = 1.0
-) -> Tuple[List[int], List[float], List[List[float]]]:
-    """
-    Slices candidate token logits and computes calibrated softmax probabilities.
-    last_logits: tensor of shape (M, vocab_size) on device.
-    """
-    inv_t = 1.0 / max(temperature, 1e-4)
-    win_indices = []
-    win_probs = []
-    all_field_probs = []
-
-    for i, cands_list in enumerate(cands_per_field):
-        field_logits = last_logits[i]
-        choice_scores = []
-        for ids in cands_list:
-            if not ids:
-                choice_scores.append(-1e9)
-            elif len(ids) == 1:
-                choice_scores.append(field_logits[ids[0]].item())
-            else:
-                choice_scores.append(field_logits[ids].max().item())
-
-        scores_t = torch.tensor(choice_scores, dtype=torch.float32) * inv_t
-        probs = F.softmax(scores_t, dim=-1).tolist()
-
-        best_idx = int(torch.argmax(scores_t).item())
-        win_indices.append(best_idx)
-        win_probs.append(probs[best_idx])
-        all_field_probs.append(probs)
-
-    return win_indices, win_probs, all_field_probs
-
-
 @gpu_decorator
 def run_parallel_generation_torch(
     context: str,
@@ -118,8 +83,11 @@ def run_parallel_generation_torch(
     # 1. Compile schema metadata
     meta = schema.compile_parallel_metadata(tokenizer)
     field_items = meta["field_items"]
+    suffix_lengths = meta["suffix_lengths"]
     cands_per_field = meta["cands_per_field"]
-    suffixes_str = [item[2] for item in field_items]
+    prefixes = meta["prefixes"]
+    has_collisions = meta["has_collisions"]
+    suffixes_batch = meta["suffixes_batch"]
     M = len(field_items)
 
     # 2. High-density semantic catalog prefill
@@ -143,18 +111,8 @@ def run_parallel_generation_torch(
     t_suf0 = time.perf_counter()
     pad_id = tokenizer.pad_token_id or tokenizer.eos_token_id or 0
 
-    suffix_tokens_list = [tokenizer.encode(s, add_special_tokens=False) for s in suffixes_str]
-    max_len = max(len(s) for s in suffix_tokens_list)
-    padded_suffixes = []
-    attention_masks = []
-
-    for s in suffix_tokens_list:
-        pad_len = max_len - len(s)
-        padded_suffixes.append([pad_id] * pad_len + s)
-        attention_masks.append([0] * pad_len + [1] * len(s))
-
-    suffix_arr = torch.tensor(padded_suffixes, dtype=torch.long, device=device)
-    suffix_mask = torch.tensor(attention_masks, dtype=torch.long, device=device)
+    suffix_arr = torch.tensor(suffixes_batch, dtype=torch.long, device=device)
+    suffix_mask = (suffix_arr != pad_id).long()
 
     # Broadcast KV cache to batch size M
     with torch.no_grad():
@@ -172,29 +130,30 @@ def run_parallel_generation_torch(
         full_mask = torch.cat([prefix_mask, suffix_mask], dim=1)
 
         out = model(suffix_arr, past_key_values=batched_cache, attention_mask=full_mask)
-        # Logits at the last token position for each field
-        last_logits = out.logits[:, -1, :]
+        suffix_out = out.logits
 
     t_suffix_eval = (time.perf_counter() - t_suf0) * 1000
 
-    # 4. Slicing & Softmax
-    win_indices, win_probs, all_field_probs = _extract_choice_logits(
-        last_logits, cands_per_field, temperature=temperature
-    )
-
-    # 5. Assemble typed JSON and telemetry
+    # 4. Slicing, Disambiguation & Softmax
     parsed_json = {}
     field_telemetry = {}
 
-    for i, (fname, fdef, _) in enumerate(field_items):
-        best_idx = win_indices[i]
-        w_prob = win_probs[i]
-        all_probs = all_field_probs[i]
+    for i, (fname, fdef) in enumerate(field_items):
+        decision_idx = suffix_lengths[i] - 1
+        field_logits = suffix_out[i, decision_idx, :]
+        cand_tokens = cands_per_field[i]
+
+        scores = [float(field_logits[tid].item()) for tid in cand_tokens]
+        scores_t = torch.tensor(scores, dtype=torch.float32) / max(temperature, 1e-4)
+        probs = F.softmax(scores_t, dim=-1).tolist()
+        w_idx = int(torch.argmax(scores_t).item())
+        w_prob = float(probs[w_idx])
+        all_probs = probs
 
         if fdef.field_type == "boolean":
-            val = (best_idx == 0)
+            val = (w_idx == 0)
         else:
-            val = fdef.choices[best_idx]
+            val = fdef.choices[w_idx]
 
         parsed_json[fname] = {
             "value": val,