symbolicai 0.21.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

symai/server/huggingface_server.py

@@ -1,39 +1,112 @@
 import argparse
 import random
-from typing import List, Optional
 
 import numpy as np
 import torch
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          BitsAndBytesConfig, StoppingCriteria,
-                          StoppingCriteriaList)
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    StoppingCriteria,
+    StoppingCriteriaList,
+)
 
 # General arguments
 parser = argparse.ArgumentParser(description="FastAPI server for Hugging Face models")
 parser.add_argument("--model", type=str, help="Path to the model")
-parser.add_argument("--host", type=str, default="localhost", help="Host address. (default: localhost)")
+parser.add_argument(
+    "--host", type=str, default="localhost", help="Host address. (default: localhost)"
+)
 parser.add_argument("--port", type=int, default=8000, help="Port number. (default: 8000)")
 
 # Quantization arguments with 'quant_' prefix
-parser.add_argument("--quant", action="store_true", default=False, help="Enable quantization; see help for available quantization options (default: False)")
-parser.add_argument("--quant_load_in_8bit", action="store_true", default=False, help="Load model in 8-bit precision (default: False)")
-parser.add_argument("--quant_load_in_4bit", action="store_true", default=False, help="Load model in 4-bit precision (default: False)")
-parser.add_argument("--quant_llm_int8_threshold", type=float, default=6.0, help="LLM int8 threshold (default: 6.0)")
-parser.add_argument("--quant_llm_int8_skip_modules", type=str, nargs="+", default=None, help="LLM int8 skip modules (default: None)")
-parser.add_argument("--quant_llm_int8_enable_fp32_cpu_offload", action="store_true", default=False, help="Enable FP32 CPU offload for LLM int8 (default: False)")
-parser.add_argument("--quant_llm_int8_has_fp16_weight", action="store_true", default=False, help="LLM int8 has FP16 weight (default: False)")
-parser.add_argument("--quant_bnb_4bit_compute_dtype", type=str, default=None, help="BNB 4-bit compute dtype (default: None)")
-parser.add_argument("--quant_bnb_4bit_quant_type", type=str, default="fp4", help="BNB 4-bit quantization type (default: fp4)")
-parser.add_argument("--quant_bnb_4bit_use_double_quant", action="store_true", default=False, help="Use double quantization for BNB 4-bit (default: False)")
-parser.add_argument("--quant_bnb_4bit_quant_storage", type=str, default=None, help="BNB 4-bit quantization storage (default: None)")
+parser.add_argument(
+    "--quant",
+    action="store_true",
+    default=False,
+    help="Enable quantization; see help for available quantization options (default: False)",
+)
+parser.add_argument(
+    "--quant_load_in_8bit",
+    action="store_true",
+    default=False,
+    help="Load model in 8-bit precision (default: False)",
+)
+parser.add_argument(
+    "--quant_load_in_4bit",
+    action="store_true",
+    default=False,
+    help="Load model in 4-bit precision (default: False)",
+)
+parser.add_argument(
+    "--quant_llm_int8_threshold", type=float, default=6.0, help="LLM int8 threshold (default: 6.0)"
+)
+parser.add_argument(
+    "--quant_llm_int8_skip_modules",
+    type=str,
+    nargs="+",
+    default=None,
+    help="LLM int8 skip modules (default: None)",
+)
+parser.add_argument(
+    "--quant_llm_int8_enable_fp32_cpu_offload",
+    action="store_true",
+    default=False,
+    help="Enable FP32 CPU offload for LLM int8 (default: False)",
+)
+parser.add_argument(
+    "--quant_llm_int8_has_fp16_weight",
+    action="store_true",
+    default=False,
+    help="LLM int8 has FP16 weight (default: False)",
+)
+parser.add_argument(
+    "--quant_bnb_4bit_compute_dtype",
+    type=str,
+    default=None,
+    help="BNB 4-bit compute dtype (default: None)",
+)
+parser.add_argument(
+    "--quant_bnb_4bit_quant_type",
+    type=str,
+    default="fp4",
+    help="BNB 4-bit quantization type (default: fp4)",
+)
+parser.add_argument(
+    "--quant_bnb_4bit_use_double_quant",
+    action="store_true",
+    default=False,
+    help="Use double quantization for BNB 4-bit (default: False)",
+)
+parser.add_argument(
+    "--quant_bnb_4bit_quant_storage",
+    type=str,
+    default=None,
+    help="BNB 4-bit quantization storage (default: None)",
+)
 
 # Model inference arguments
 # https://huggingface.co/docs/transformers/main/en/main_classes/model
-parser.add_argument("--torch_dtype", type=str, default="auto", help="A string that is a valid torch.dtype. E.g. “float32” loads the model in torch.float32, “float16” loads in torch.float16 etc. (default: auto)")
-parser.add_argument("--attn_implementation", type=str, default="sdpa", help='The attention implementation to use in the model (if relevant). Can be any of "eager" (manual implementation of the attention), "sdpa" (using F.scaled_dot_product_attention), or "flash_attention_2" (using Dao-AILab/flash-attention). By default, if available, SDPA will be used for torch>=2.1.1. (default: sdpa)')
-parser.add_argument("--device_map", type=str, default="auto", help="A string that is a valid device. E.g. “cuda” loads the model on the GPU, “cpu” loads it on the CPU. (default: auto)")
+parser.add_argument(
+    "--torch_dtype",
+    type=str,
+    default="auto",
+    help="A string that is a valid torch.dtype. E.g. “float32” loads the model in torch.float32, “float16” loads in torch.float16 etc. (default: auto)",
+)
+parser.add_argument(
+    "--attn_implementation",
+    type=str,
+    default="sdpa",
+    help='The attention implementation to use in the model (if relevant). Can be any of "eager" (manual implementation of the attention), "sdpa" (using F.scaled_dot_product_attention), or "flash_attention_2" (using Dao-AILab/flash-attention). By default, if available, SDPA will be used for torch>=2.1.1. (default: sdpa)',
+)
+parser.add_argument(
+    "--device_map",
+    type=str,
+    default="auto",
+    help="A string that is a valid device. E.g. “cuda” loads the model on the GPU, “cpu” loads it on the CPU. (default: auto)",
+)
 
 args = parser.parse_args()
 
@@ -47,7 +120,7 @@ quant_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=args.quant_bnb_4bit_compute_dtype,
     bnb_4bit_quant_type=args.quant_bnb_4bit_quant_type,
     bnb_4bit_use_double_quant=args.quant_bnb_4bit_use_double_quant,
-    bnb_4bit_quant_storage=args.quant_bnb_4bit_quant_storage
+    bnb_4bit_quant_storage=args.quant_bnb_4bit_quant_storage,
 )
 
 app = FastAPI()
@@ -70,6 +143,7 @@ else:
         attn_implementation=args.attn_implementation,
     )
 
+
 def set_seed(seed):
     random.seed(seed)
     np.random.seed(seed)
@@ -77,6 +151,7 @@ def set_seed(seed):
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(seed)
 
+
 class StoppingCriteriaSub(StoppingCriteria):
     # https://discuss.huggingface.co/t/implimentation-of-stopping-criteria-list/20040/13
     def __init__(self, stop_words, tokenizer):
@@ -84,39 +159,48 @@ class StoppingCriteriaSub(StoppingCriteria):
         self.stop_words = stop_words
         self.tokenizer = tokenizer
 
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+    def __call__(self, input_ids: torch.LongTensor, _scores: torch.FloatTensor, **_kwargs) -> bool:
         for stop_word in self.stop_words:
-            if self.tokenizer.decode(input_ids[0][-len(self.tokenizer.encode(stop_word)):]).strip() == stop_word:
+            if (
+                self.tokenizer.decode(
+                    input_ids[0][-len(self.tokenizer.encode(stop_word)) :]
+                ).strip()
+                == stop_word
+            ):
                 return True
         return False
 
+
 class TokenizeRequest(BaseModel):
     input: str
-    add_special_tokens: Optional[bool] = False
+    add_special_tokens: bool | None = False
+
 
 class DetokenizeRequest(BaseModel):
-    tokens: List[int]
-    skip_special_tokens: Optional[bool] = True
+    tokens: list[int]
+    skip_special_tokens: bool | None = True
+
 
 class ChatCompletionRequest(BaseModel):
-    messages: List[dict]
-    temperature: float = 1.
-    top_p: float = 1.
-    stop: Optional[List[str]] = None
-    seed: Optional[int] = None
-    max_tokens: Optional[int] = 2048
-    max_tokens_forcing: Optional[int] = None
+    messages: list[dict]
+    temperature: float = 1.0
+    top_p: float = 1.0
+    stop: list[str] | None = None
+    seed: int | None = None
+    max_tokens: int | None = 2048
+    max_tokens_forcing: int | None = None
     top_k: int = 50
     logprobs: bool = False
     do_sample: bool = True
     num_beams: int = 1
     num_beam_groups: int = 1
-    eos_token_id: Optional[int] = None
+    eos_token_id: int | None = None
+
 
 @app.post("/chat")
 def chat_completions(request: ChatCompletionRequest):
     chat = request.messages
-    #@TODO: is there a way to assert that the loaded model has chat capabilities?
+    # @TODO: is there a way to assert that the loaded model has chat capabilities?
     inputs = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(inputs, return_tensors="pt").to(model.device)
 
@@ -129,30 +213,29 @@ def chat_completions(request: ChatCompletionRequest):
         "do_sample": request.do_sample,
         "num_beams": request.num_beams,
         "num_beam_groups": request.num_beam_groups,
-        "eos_token_id": request.eos_token_id if request.eos_token_id is not None else tokenizer.eos_token_id,
+        "eos_token_id": request.eos_token_id
+        if request.eos_token_id is not None
+        else tokenizer.eos_token_id,
         "output_logits": request.logprobs,
         "return_dict_in_generate": True,
     }
 
     if request.stop:
-        generation_config["stopping_criteria"] = StoppingCriteriaList([StoppingCriteriaSub(stop_words=request.stop, tokenizer=tokenizer)])
+        generation_config["stopping_criteria"] = StoppingCriteriaList(
+            [StoppingCriteriaSub(stop_words=request.stop, tokenizer=tokenizer)]
+        )
 
     if request.seed:
         set_seed(request.seed)
 
     outputs = model.generate(**inputs, **generation_config)
 
-    new_tokens = outputs.sequences[0][inputs.input_ids.shape[-1]:]
+    new_tokens = outputs.sequences[0][inputs.input_ids.shape[-1] :]
     generated_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
 
-    response = {
+    return {
         "choices": [
-            {
-                "message": {
-                    "role": "assistant",
-                    "content": generated_text
-                }
-            },
+            {"message": {"role": "assistant", "content": generated_text}},
         ],
         "metadata": {
             "model": args.model,
@@ -172,23 +255,26 @@ def chat_completions(request: ChatCompletionRequest):
             "logits": [logits.tolist() for logits in outputs.logits],
             "model_input": chat,
             "model_chat_format": tokenizer.chat_template,
-        }
+        },
     }
 
-    return response
 
 @app.post("/tokenize")
 def tokenize(request: TokenizeRequest):
     tokens = tokenizer.encode(request.input, add_special_tokens=request.add_special_tokens)
     return {"tokens": tokens}
 
+
 @app.post("/detokenize")
 def detokenize(request: DetokenizeRequest):
     text = tokenizer.decode(request.tokens, skip_special_tokens=request.skip_special_tokens)
     return {"text": text}
 
+
 def huggingface_server():
-    import uvicorn
-    from functools import partial
+    # Lazy imports keep optional server dependencies out of module import path.
+    from functools import partial  # noqa
+    import uvicorn  # noqa
+
     command = partial(uvicorn.run, app)
     return command, args

symai/server/llama_cpp_server.py

@@ -1,45 +1,55 @@
 import argparse
 import subprocess
 import sys
-import os
+from pathlib import Path
+
 from loguru import logger
 
+
 def llama_cpp_server():
     parser = argparse.ArgumentParser(description="A wrapper for llama_cpp.", add_help=False)
-    parser.add_argument("--help", action="store_true", help="Show available options for llama_cpp server.")
-    parser.add_argument("--env", choices=['python', 'cpp'], default='python',
-                      help="Choose programming environment (python or cpp)")
+    parser.add_argument(
+        "--help", action="store_true", help="Show available options for llama_cpp server."
+    )
+    parser.add_argument(
+        "--env",
+        choices=["python", "cpp"],
+        default="python",
+        help="Choose programming environment (python or cpp)",
+    )
     parser.add_argument("--cpp-server-path", type=str, help="Path to llama.cpp server executable")
 
     main_args, llama_cpp_args = parser.parse_known_args()
 
     if main_args.help:
-        if main_args.env == 'python':
+        if main_args.env == "python":
             command = [sys.executable, "-m", "llama_cpp.server", "--help"]
-            subprocess.run(command)
+            subprocess.run(command, check=False)
         else:
             if not main_args.cpp_server_path:
                 logger.error("Error: --cpp-server-path is required when using cpp environment")
                 sys.exit(1)
-            if not os.path.exists(main_args.cpp_server_path):
+            if not Path(main_args.cpp_server_path).exists():
                 logger.error(f"Error: Executable not found at {main_args.cpp_server_path}")
                 sys.exit(1)
             command = [main_args.cpp_server_path, "--help"]
-            subprocess.run(command)
+            subprocess.run(command, check=False)
         sys.exit(0)
 
-    if main_args.env == 'cpp':
+    if main_args.env == "cpp":
         if not main_args.cpp_server_path:
             logger.error("Error: --cpp-server-path is required when using cpp environment")
             sys.exit(1)
-        if not os.path.exists(main_args.cpp_server_path):
+        if not Path(main_args.cpp_server_path).exists():
             logger.error(f"Error: Executable not found at {main_args.cpp_server_path}")
             sys.exit(1)
         command = [
             main_args.cpp_server_path,
             *llama_cpp_args,
         ]
-        llama_cpp_args = [arg for arg in llama_cpp_args if not arg.startswith("--embedding")] # Exclude embedding argument
+        llama_cpp_args = [
+            arg for arg in llama_cpp_args if not arg.startswith("--embedding")
+        ]  # Exclude embedding argument
     else:  # python
         command = [
             sys.executable,

symai/server/qdrant_server.py

@@ -0,0 +1,206 @@
+import argparse
+import subprocess
+import sys
+from pathlib import Path
+
+from loguru import logger
+
+
+def qdrant_server():  # noqa
+    """
+    A wrapper for Qdrant server that supports both Docker and binary execution modes.
+
+    Returns:
+        tuple: (command, args) where command is the list to execute and args are the parsed arguments
+    """
+    parser = argparse.ArgumentParser(description="A wrapper for Qdrant server.", add_help=False)
+    parser.add_argument(
+        "--help", action="store_true", help="Show available options for Qdrant server."
+    )
+    parser.add_argument(
+        "--env",
+        choices=["docker", "binary"],
+        default="docker",
+        help="Choose execution environment (docker or binary)",
+    )
+    parser.add_argument("--binary-path", type=str, help="Path to Qdrant binary executable")
+    parser.add_argument(
+        "--docker-image",
+        type=str,
+        default="qdrant/qdrant:latest",
+        help="Docker image to use (default: qdrant/qdrant:latest)",
+    )
+    parser.add_argument(
+        "--host", type=str, default="0.0.0.0", help="Host address to bind to (default: 0.0.0.0)"
+    )
+    parser.add_argument("--port", type=int, default=6333, help="REST API port (default: 6333)")
+    parser.add_argument("--grpc-port", type=int, default=6334, help="gRPC API port (default: 6334)")
+    parser.add_argument(
+        "--storage-path",
+        type=str,
+        default="./qdrant_storage",
+        help="Path to Qdrant storage directory (default: ./qdrant_storage)",
+    )
+    parser.add_argument(
+        "--config-path", type=str, default=None, help="Path to Qdrant configuration file"
+    )
+    parser.add_argument(
+        "--docker-container-name",
+        type=str,
+        default="qdrant",
+        help="Name for Docker container (default: qdrant)",
+    )
+    parser.add_argument(
+        "--docker-remove",
+        action="store_true",
+        default=True,
+        help="Remove container when it stops (default: True)",
+    )
+    parser.add_argument(
+        "--docker-detach",
+        action="store_true",
+        default=False,
+        help="Run Docker container in detached mode (default: False)",
+    )
+
+    main_args, qdrant_args = parser.parse_known_args()
+
+    if main_args.help:
+        if main_args.env == "docker":
+            # Show Docker help
+            command = ["docker", "run", "--rm", main_args.docker_image, "--help"]
+            subprocess.run(command, check=False)
+        else:
+            if not main_args.binary_path:
+                logger.error("Error: --binary-path is required when using binary environment")
+                sys.exit(1)
+            if not Path(main_args.binary_path).exists():
+                logger.error(f"Error: Binary not found at {main_args.binary_path}")
+                sys.exit(1)
+            command = [main_args.binary_path, "--help"]
+            subprocess.run(command, check=False)
+        sys.exit(0)
+
+    if main_args.env == "binary":
+        if not main_args.binary_path:
+            logger.error("Error: --binary-path is required when using binary environment")
+            sys.exit(1)
+        if not Path(main_args.binary_path).exists():
+            logger.error(f"Error: Binary not found at {main_args.binary_path}")
+            sys.exit(1)
+
+        # Build command for binary execution
+        command = [main_args.binary_path]
+
+        # Ensure storage directory exists
+        storage_path = Path(main_args.storage_path)
+        storage_path.mkdir(parents=True, exist_ok=True)
+        abs_storage_path = str(storage_path.resolve())
+
+        # Add standard Qdrant arguments
+        # Set storage path via environment variable or command argument
+        # Qdrant binary accepts --storage-path argument
+        command.extend(["--storage-path", abs_storage_path])
+
+        # Add host, port, and grpc-port arguments
+        command.extend(["--host", main_args.host])
+        command.extend(["--port", str(main_args.port)])
+        command.extend(["--grpc-port", str(main_args.grpc_port)])
+
+        if main_args.config_path:
+            command.extend(["--config-path", main_args.config_path])
+
+        # Add any additional Qdrant-specific arguments
+        command.extend(qdrant_args)
+
+    else:  # docker
+        # Ensure storage directory exists
+        storage_path = Path(main_args.storage_path)
+        storage_path.mkdir(parents=True, exist_ok=True)
+        abs_storage_path = str(storage_path.resolve())
+
+        # Build Docker command
+        command = ["docker", "run"]
+
+        # Container management options
+        if main_args.docker_remove:
+            command.append("--rm")
+
+        if main_args.docker_detach:
+            command.append("-d")
+        # Note: We don't add -it by default to avoid issues in non-interactive environments
+        # Users can add it manually if needed via qdrant_args
+
+        # Container name
+        command.extend(["--name", main_args.docker_container_name])
+
+        # Port mappings
+        command.extend(["-p", f"{main_args.port}:6333"])
+        command.extend(["-p", f"{main_args.grpc_port}:6334"])
+
+        # Volume mount for storage
+        command.extend(["-v", f"{abs_storage_path}:/qdrant/storage:z"])
+
+        # Volume mount for config (if provided)
+        # Note: Qdrant Docker image accepts environment variables and config files
+        # For custom config, mount it as a volume before the image name
+        if main_args.config_path:
+            config_path = Path(main_args.config_path)
+            abs_config_path = config_path.resolve()
+            config_dir = str(abs_config_path.parent)
+            command.extend(["-v", f"{config_dir}:/qdrant/config:z"])
+            # Qdrant looks for config.yaml in /qdrant/config by default
+
+        # Set storage path environment variable to use the mounted volume
+        command.extend(["-e", "QDRANT__STORAGE__STORAGE_PATH=/qdrant/storage"])
+
+        # Docker image
+        command.append(main_args.docker_image)
+
+        # Qdrant server arguments (if any additional ones are passed)
+
+        # Add any additional Qdrant arguments
+        if qdrant_args:
+            command.extend(qdrant_args)
+
+    # Prepare args for config storage (similar to llama_cpp_server pattern)
+    # Extract key-value pairs for configuration
+    config_args = []
+    if main_args.env == "docker":
+        config_args = [
+            "--env",
+            main_args.env,
+            "--host",
+            main_args.host,
+            "--port",
+            str(main_args.port),
+            "--grpc-port",
+            str(main_args.grpc_port),
+            "--storage-path",
+            main_args.storage_path,
+            "--docker-image",
+            main_args.docker_image,
+            "--docker-container-name",
+            main_args.docker_container_name,
+        ]
+        if main_args.config_path:
+            config_args.extend(["--config-path", main_args.config_path])
+    else:
+        config_args = [
+            "--env",
+            main_args.env,
+            "--binary-path",
+            main_args.binary_path,
+            "--host",
+            main_args.host,
+            "--port",
+            str(main_args.port),
+            "--grpc-port",
+            str(main_args.grpc_port),
+            "--storage-path",
+            main_args.storage_path,
+        ]
+        if main_args.config_path:
+            config_args.extend(["--config-path", main_args.config_path])
+
+    return command, config_args