Format code and fix missing arguments (#710)

Author: zyzshishui

examples/search-r1/generate_with_search.py

@@ -16,28 +16,23 @@
     "max_turns": 2,
     "topk": 3,
     "search_concurrency": 256,
-
     # ============== Search Backend Selection ==============
     "search_backend": "local",  # Options: "local" or "google"
-
     # ============== Local Search Configuration ==============
     # (Only used when search_backend="local")
     "local": {
         "search_url": "http://127.0.0.1:8000/retrieve",  # URL of your local retrieval server
         "proxy": None,  # Set to your proxy if needed
     },
-
     # ============== Google Search Configuration ==============
     # (Only used when search_backend="google")
     "google": {
         "api_key": "your_api_key_here",  # Replace with your actual API key
         "snippet_only": True,  # Set to True to only return snippets
         "proxy": None,  # Set to your proxy if needed
     },
-
     # ============== Log Probability Collection ==============
     "return_logprob": True,  # Set to True to collect log probabilities for TIS metrics
-
     # ============== Reward Model Configuration ==============
     "format_score": 0.2,
 }
@@ -90,10 +85,7 @@ async def search(query: str) -> str:
             proxy=google_config["proxy"],
         )
     else:
-        raise ValueError(
-            f"Unknown search backend: {backend}. "
-            f"Must be either 'local' or 'google'."
-        )
+        raise ValueError(f"Unknown search backend: {backend}. " f"Must be either 'local' or 'google'.")
 
     return _passages2string(result)
 

examples/search-r1/local_dense_retriever/retrieval_server.py

@@ -355,9 +355,7 @@ def retrieve_endpoint(request: QueryRequest):
         request.topk = config.retrieval_topk  # fallback to default
 
     # Perform batch retrieval
-    tmp = retriever.batch_search(
-        query_list=request.queries, num=request.topk, return_score=request.return_scores
-    )
+    tmp = retriever.batch_search(query_list=request.queries, num=request.topk, return_score=request.return_scores)
 
     scores = []
     try:
@@ -382,7 +380,10 @@ def retrieve_endpoint(request: QueryRequest):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Launch the local faiss retriever.")
     parser.add_argument(
-        "--index_path", type=str, default="/home/peterjin/mnt/index/wiki-18/e5_Flat.index", help="Corpus indexing file."
+        "--index_path",
+        type=str,
+        default="/home/peterjin/mnt/index/wiki-18/e5_Flat.index",
+        help="Corpus indexing file.",
     )
     parser.add_argument(
         "--corpus_path",

examples/search-r1/local_search_server.py

@@ -19,9 +19,8 @@
         }
 """
 
-import asyncio
-import argparse
-from typing import List, Dict, Optional
+from typing import Dict, List, Optional
+
 import aiohttp
 
 
@@ -55,7 +54,7 @@ async def local_search(
     payload = {
         "queries": [query],
         "topk": top_k,
-        "return_scores": False  # We don't need scores for compatibility with google_search_server
+        "return_scores": False,  # We don't need scores for compatibility with google_search_server
     }
 
     # Send async request to local retrieval server
@@ -86,19 +85,15 @@ async def local_search(
         # retrieval_server returns: {"document": {"id": "...", "contents": '"Title"\nText...'}}
         if isinstance(item, dict):
             # Access the document dict first, then get contents
-            content = item.get('contents', "")
+            content = item.get("contents", "")
 
             if content:
                 # The contents are already in the correct format: '"Title"\nText content...'
                 # Just pass through as-is to match google_search format
-                contexts.append({
-                    "document": {"contents": content}
-                })
+                contexts.append({"document": {"contents": content}})
             else:
                 # Empty content case - provide default values
-                contexts.append({
-                    "document": {"contents": '"No title."\nNo snippet available.'}
-                })
+                contexts.append({"document": {"contents": '"No title."\nNo snippet available.'}})
 
     # If no results found, return empty list (consistent with google_search_server.py)
     return contexts

slime/backends/fsdp_utils/arguments.py

@@ -28,6 +28,7 @@ class FSDPArgs:
 
     # Precision
     gradient_checkpointing: bool = False
+    fp16: bool = False
 
     # FSDP configuration
     fsdp_full_params: bool = False  # If True, use full_tensor; if False, use shard_tensor

slime/backends/megatron_utils/loss.py

@@ -725,4 +725,4 @@ def loss_function(
                 device=logits.device,
             ),
         },
-    )
+    )