fix: support messages field in calibrate_draft_vocab and compute_hidden_states

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenhan Yu <chenhany@nvidia.com>

Author: ChenhanYu

examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py

@@ -201,7 +201,7 @@ async def submit_generates():
         for entry in dataset:
             conversation_id = entry.get("conversation_id", entry.get("uuid"))
 
-            conversations = entry["conversations"]
+            conversations = entry.get("messages") or entry["conversations"]
             if not conversations or not isinstance(conversations, list):
                 num_invalid += 1
                 continue

examples/speculative_decoding/scripts/calibrate_draft_vocab.py

@@ -49,7 +49,9 @@ def main():
     tokenizer = AutoTokenizer.from_pretrained(args.model)
     with open(args.data) as f:
         lines = islice(f, args.calibrate_size) if args.calibrate_size else f
-        conversations = [json.loads(line)["conversations"] for line in lines]
+        conversations = [
+            (d := json.loads(line)).get("messages") or d["conversations"] for line in lines
+        ]
         conversations = [item for sublist in conversations for item in sublist]
 
     d2t = calibrate_frequent_vocab(tokenizer, conversations, args.draft_vocab_size)