Remove OfflineEagleRegistry; Benjamin:Fix conversation ID indexing

Signed-off-by: h-guo18 <[email protected]>

Author: h-guo18

examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py

@@ -84,10 +84,6 @@ async def main(args: argparse.Namespace) -> None:
     with args.input_file.open("r", encoding="utf-8") as f:
         all_conversations.extend([json.loads(line) for line in f if line.strip()])
 
-    if any(not entry.get("conversation_id") for entry in all_conversations):
-        msg = "All conversations must have a 'conversation_id' field."
-        raise ValueError(msg)
-
     print("Loaded", len(all_conversations), "conversations from", args.input_file)
 
     model = AutoModel.from_pretrained(args.model, torch_dtype="auto", device_map="auto")

examples/speculative_decoding/collect_hidden_states/send_conversations_for_hiddens.py

@@ -95,10 +95,6 @@ async def main(args: argparse.Namespace) -> None:
     with args.input_file.open("r", encoding="utf-8") as f:
         all_conversations.extend([json.loads(line) for line in f if line.strip()])
 
-    if any(not entry.get("conversation_id") for entry in all_conversations):
-        msg = "All conversations must have a 'conversation_id' field."
-        raise ValueError(msg)
-
     print("Loaded", len(all_conversations), "conversations from", args.input_file)
 
     client: AsyncOpenAI = AsyncOpenAI(
@@ -127,12 +123,14 @@ async def main(args: argparse.Namespace) -> None:
     num_total_conversations = min(
         len(all_conversations), args.debug_max_num_conversations or len(all_conversations)
     )
-    for entry in tqdm(
-        all_conversations[: args.debug_max_num_conversations],
-        desc="Processing conversations",
-        total=num_total_conversations,
+    for idx, entry in enumerate(
+        tqdm(
+            all_conversations[: args.debug_max_num_conversations],
+            desc="Processing conversations",
+            total=num_total_conversations,
+        )
     ):
-        conversation_id = entry["conversation_id"]
+        conversation_id = entry.get("conversation_id", "{:08d}".format(idx))
         conversations = entry["conversations"]
         if not conversations or not isinstance(conversations, list):
             num_invalid += 1

examples/speculative_decoding/eagle_utils.py

@@ -269,12 +269,12 @@ def make_eagle_supervised_data_module(
 
         # Filter to conversations that exist in the offline data and in the provided json
         valid_entries = []
-        for entry in data_json:
-            conv_id = entry.get("conversation_id") or entry.get("id")
-            if not conv_id:
-                raise ValueError(
-                    "Each entry in the data json must have a 'conversation_id' or 'id' field."
-                )
+        for idx, entry in enumerate(data_json):
+            conv_id = entry.get("conversation_id")
+            if conv_id is None:
+                conv_id = entry.get("id")
+            if conv_id is None:
+                conv_id = "{:08d}".format(idx)
             file_path = str(offline_data_path / f"{conv_id}.pt")
             if file_path in all_files:
                 valid_entries.append((entry, file_path))

modelopt/torch/speculative/plugins/transformers.py

@@ -330,7 +330,6 @@ def forward(
 
 
 @EagleDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"})
-@OfflineEagleDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"})
 class HFEagleModel(EagleModel):
     """Eagle Model Class for huggingface models."""