refactor: 优化工具脚本接口并修复批处理问题

scripts/tools/server.py

@@ -1,6 +1,8 @@
 import argparse
 from pathlib import Path
 
+import torch
+
 from astrai.inference.server import run_server
 
 
@@ -21,17 +23,49 @@ def main():
         default=None,
         help="Path to model parameters (default: project_root/params)",
     )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="Device to load model on (default: cuda)",
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16",
+        choices=["bfloat16", "float16", "float32"],
+        help="Data type for model weights (default: bfloat16)",
+    )
+    parser.add_argument(
+        "--max_batch_size",
+        type=int,
+        default=16,
+        help="Maximum batch size for continuous batching (default: 16)",
+    )
     args = parser.parse_args()
 
-    # If param_path is provided, set environment variable or modify global?
-    # Currently the server loads from default location on startup.
-    # We could pass it via an environment variable, but for simplicity we assume
-    # the default location is correct.
-    project_root = Path(__file__).parent.parent
+    # Convert dtype string to torch dtype
+    dtype_map = {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }
+    dtype = dtype_map[args.dtype]
+
+    project_root = Path(__file__).parent.parent.parent
     param_path = args.param_path or (project_root / "params")
     print(f"Starting AstrAI inference server on http://{args.host}:{args.port}")
-    print(f"Model parameters expected at: {[param_path]}")
-    run_server(host=args.host, port=args.port, reload=args.reload)
+    print(f"Model parameters expected at: {param_path}")
+    print(f"Device: {args.device}, Dtype: {args.dtype}")
+    run_server(
+        host=args.host,
+        port=args.port,
+        reload=args.reload,
+        device=args.device,
+        dtype=dtype,
+        param_path=param_path,
+        max_batch_size=args.max_batch_size,
+    )
 
 
 if __name__ == "__main__":