fix pre-commit

Author: zyzshishui

slime/backends/megatron_utils/data.py

@@ -174,7 +174,15 @@ def log_rollout_data(rollout_id, args, rollout_data):
         log_dict = {}
         response_lengths = rollout_data["response_lengths"]
         for key, val in rollout_data.items():
-            if key == "tokens" or key == "loss_masks" or key == "sample_indices"or key == "rollout_time" or key == "completion_tokens_stats" or key == "partial_samples" or key == "total_off_policy_tokens":
+            if (
+                key == "tokens"
+                or key == "loss_masks"
+                or key == "sample_indices"
+                or key == "rollout_time"
+                or key == "completion_tokens_stats"
+                or key == "partial_samples"
+                or key == "total_off_policy_tokens"
+            ):
                 continue
             # Upload per sample mean for each rollout value
             # There are the following assumptions:
@@ -248,7 +256,9 @@ def log_partial_rollout_data(rollout_id, args, rollout_data):
         total_off_policy_tokens = rollout_data["total_off_policy_tokens"]
         if total_off_policy_tokens is not None:
             log_dict["total_off_policy_tokens"] = total_off_policy_tokens
-            log_dict["off_policy_ratio"] = total_off_policy_tokens / (log_dict["total_tokens"] + total_off_policy_tokens)
+            log_dict["off_policy_ratio"] = total_off_policy_tokens / (
+                log_dict["total_tokens"] + total_off_policy_tokens
+            )
 
         response_lengths = rollout_data["response_lengths"]
 
@@ -298,7 +308,7 @@ def log_partial_rollout_data(rollout_id, args, rollout_data):
                 dst=mpu.get_data_parallel_src_rank(with_context_parallel=True),
                 group=mpu.get_data_parallel_group(with_context_parallel=True),
             )
-            
+
 
 def log_multi_turn_data(rollout_id, args, rollout_data):
     if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():

slime/rollout/sglang_example.py

@@ -288,19 +288,21 @@ async def generate_rollout_async(args, rollout_id: int, data_source) -> list[lis
     if state.completion_tokens_list:
         completion_tokens_array = np.array(state.completion_tokens_list)
         completion_tokens_stats = {
-            'total_completion_tokens': np.sum(completion_tokens_array).item(),
-            'completion_tokens_mean': np.mean(completion_tokens_array).item(),
-            'completion_tokens_std': np.std(completion_tokens_array).item(),
-            'completion_tokens_count': len(completion_tokens_array),
+            "total_completion_tokens": np.sum(completion_tokens_array).item(),
+            "completion_tokens_mean": np.mean(completion_tokens_array).item(),
+            "completion_tokens_std": np.std(completion_tokens_array).item(),
+            "completion_tokens_count": len(completion_tokens_array),
         }
 
     if len(data) > 0:
-        data[0][0].metadata.update({
-            'rollout_time': rollout_time,
-            'completion_tokens_stats': completion_tokens_stats,
-            'partial_samples': state.partial_samples_count,
-            'total_off_policy_tokens': state.total_off_policy_tokens,
-        })
+        data[0][0].metadata.update(
+            {
+                "rollout_time": rollout_time,
+                "completion_tokens_stats": completion_tokens_stats,
+                "partial_samples": state.partial_samples_count,
+                "total_off_policy_tokens": state.total_off_policy_tokens,
+            }
+        )
     if completion_tokens_stats:
         print(f"[DEBUG] Rollout {rollout_id}: Completion tokens stats: {completion_tokens_stats}", flush=True)
 

slime/utils/types.py

@@ -20,7 +20,7 @@ class Sample:
     reward: Optional[Union[float, dict[str, Any]]] = None
     loss_mask: Optional[list[int]] = None
     completion_tokens: Optional[int] = None
-    
+
     class Status(Enum):
         PENDING = "pending"
         COMPLETED = "completed"