llama-stack 0.4.3__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

llama_stack-0.5.0rc1.dist-info/top_level.txt

@@ -0,0 +1 @@
+llama_stack

llama_stack/distributions/meta-reference-gpu/__init__.py

@@ -1,7 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from .meta_reference import get_distribution_template  # noqa: F401

llama_stack/distributions/meta-reference-gpu/config.yaml

@@ -1,140 +0,0 @@
-version: 2
-image_name: meta-reference-gpu
-apis:
-- agents
-- datasetio
-- eval
-- inference
-- safety
-- scoring
-- tool_runtime
-- vector_io
-providers:
-  inference:
-  - provider_id: meta-reference-inference
-    provider_type: inline::meta-reference
-    config:
-      model: ${env.INFERENCE_MODEL}
-      checkpoint_dir: ${env.INFERENCE_CHECKPOINT_DIR:=null}
-      quantization:
-        type: ${env.QUANTIZATION_TYPE:=bf16}
-      model_parallel_size: ${env.MODEL_PARALLEL_SIZE:=0}
-      max_batch_size: ${env.MAX_BATCH_SIZE:=1}
-      max_seq_len: ${env.MAX_SEQ_LEN:=4096}
-  - provider_id: sentence-transformers
-    provider_type: inline::sentence-transformers
-  vector_io:
-  - provider_id: faiss
-    provider_type: inline::faiss
-    config:
-      persistence:
-        namespace: vector_io::faiss
-        backend: kv_default
-  safety:
-  - provider_id: llama-guard
-    provider_type: inline::llama-guard
-    config:
-      excluded_categories: []
-  agents:
-  - provider_id: meta-reference
-    provider_type: inline::meta-reference
-    config:
-      persistence:
-        agent_state:
-          namespace: agents
-          backend: kv_default
-        responses:
-          table_name: responses
-          backend: sql_default
-          max_write_queue_size: 10000
-          num_writers: 4
-  eval:
-  - provider_id: meta-reference
-    provider_type: inline::meta-reference
-    config:
-      kvstore:
-        namespace: eval
-        backend: kv_default
-  datasetio:
-  - provider_id: huggingface
-    provider_type: remote::huggingface
-    config:
-      kvstore:
-        namespace: datasetio::huggingface
-        backend: kv_default
-  - provider_id: localfs
-    provider_type: inline::localfs
-    config:
-      kvstore:
-        namespace: datasetio::localfs
-        backend: kv_default
-  scoring:
-  - provider_id: basic
-    provider_type: inline::basic
-  - provider_id: llm-as-judge
-    provider_type: inline::llm-as-judge
-  - provider_id: braintrust
-    provider_type: inline::braintrust
-    config:
-      openai_api_key: ${env.OPENAI_API_KEY:=}
-  tool_runtime:
-  - provider_id: brave-search
-    provider_type: remote::brave-search
-    config:
-      api_key: ${env.BRAVE_SEARCH_API_KEY:=}
-      max_results: 3
-  - provider_id: tavily-search
-    provider_type: remote::tavily-search
-    config:
-      api_key: ${env.TAVILY_SEARCH_API_KEY:=}
-      max_results: 3
-  - provider_id: rag-runtime
-    provider_type: inline::rag-runtime
-  - provider_id: model-context-protocol
-    provider_type: remote::model-context-protocol
-storage:
-  backends:
-    kv_default:
-      type: kv_sqlite
-      db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db
-    sql_default:
-      type: sql_sqlite
-      db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/sql_store.db
-  stores:
-    metadata:
-      namespace: registry
-      backend: kv_default
-    inference:
-      table_name: inference_store
-      backend: sql_default
-      max_write_queue_size: 10000
-      num_writers: 4
-    conversations:
-      table_name: openai_conversations
-      backend: sql_default
-    prompts:
-      namespace: prompts
-      backend: kv_default
-registered_resources:
-  models:
-  - metadata: {}
-    model_id: ${env.INFERENCE_MODEL}
-    provider_id: meta-reference-inference
-    model_type: llm
-  - metadata:
-      embedding_dimension: 768
-    model_id: nomic-embed-text-v1.5
-    provider_id: sentence-transformers
-    model_type: embedding
-  shields: []
-  vector_dbs: []
-  datasets: []
-  scoring_fns: []
-  benchmarks: []
-  tool_groups:
-  - toolgroup_id: builtin::websearch
-    provider_id: tavily-search
-  - toolgroup_id: builtin::rag
-    provider_id: rag-runtime
-server:
-  port: 8321

llama_stack/distributions/meta-reference-gpu/meta_reference.py

@@ -1,163 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from pathlib import Path
-
-from llama_stack.core.datatypes import (
-    BuildProvider,
-    ModelInput,
-    Provider,
-    ShieldInput,
-    ToolGroupInput,
-)
-from llama_stack.distributions.template import DistributionTemplate, RunConfigSettings
-from llama_stack.providers.inline.inference.meta_reference import (
-    MetaReferenceInferenceConfig,
-)
-from llama_stack.providers.inline.inference.sentence_transformers import (
-    SentenceTransformersInferenceConfig,
-)
-from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
-from llama_stack_api import ModelType
-
-
-def get_distribution_template() -> DistributionTemplate:
-    providers = {
-        "inference": [BuildProvider(provider_type="inline::meta-reference")],
-        "vector_io": [
-            BuildProvider(provider_type="inline::faiss"),
-            BuildProvider(provider_type="remote::chromadb"),
-            BuildProvider(provider_type="remote::pgvector"),
-        ],
-        "safety": [BuildProvider(provider_type="inline::llama-guard")],
-        "agents": [BuildProvider(provider_type="inline::meta-reference")],
-        "eval": [BuildProvider(provider_type="inline::meta-reference")],
-        "datasetio": [
-            BuildProvider(provider_type="remote::huggingface"),
-            BuildProvider(provider_type="inline::localfs"),
-        ],
-        "scoring": [
-            BuildProvider(provider_type="inline::basic"),
-            BuildProvider(provider_type="inline::llm-as-judge"),
-            BuildProvider(provider_type="inline::braintrust"),
-        ],
-        "tool_runtime": [
-            BuildProvider(provider_type="remote::brave-search"),
-            BuildProvider(provider_type="remote::tavily-search"),
-            BuildProvider(provider_type="inline::rag-runtime"),
-            BuildProvider(provider_type="remote::model-context-protocol"),
-        ],
-    }
-    name = "meta-reference-gpu"
-    inference_provider = Provider(
-        provider_id="meta-reference-inference",
-        provider_type="inline::meta-reference",
-        config=MetaReferenceInferenceConfig.sample_run_config(
-            model="${env.INFERENCE_MODEL}",
-            checkpoint_dir="${env.INFERENCE_CHECKPOINT_DIR:=null}",
-        ),
-    )
-    embedding_provider = Provider(
-        provider_id="sentence-transformers",
-        provider_type="inline::sentence-transformers",
-        config=SentenceTransformersInferenceConfig.sample_run_config(),
-    )
-    vector_io_provider = Provider(
-        provider_id="faiss",
-        provider_type="inline::faiss",
-        config=FaissVectorIOConfig.sample_run_config(f"~/.llama/distributions/{name}"),
-    )
-
-    inference_model = ModelInput(
-        model_id="${env.INFERENCE_MODEL}",
-        provider_id="meta-reference-inference",
-    )
-    embedding_model = ModelInput(
-        model_id="nomic-embed-text-v1.5",
-        provider_id="sentence-transformers",
-        model_type=ModelType.embedding,
-        metadata={
-            "embedding_dimension": 768,
-        },
-    )
-    safety_model = ModelInput(
-        model_id="${env.SAFETY_MODEL}",
-        provider_id="meta-reference-safety",
-    )
-    default_tool_groups = [
-        ToolGroupInput(
-            toolgroup_id="builtin::websearch",
-            provider_id="tavily-search",
-        ),
-        ToolGroupInput(
-            toolgroup_id="builtin::rag",
-            provider_id="rag-runtime",
-        ),
-    ]
-
-    return DistributionTemplate(
-        name=name,
-        distro_type="self_hosted",
-        description="Use Meta Reference for running LLM inference",
-        template_path=Path(__file__).parent / "doc_template.md",
-        providers=providers,
-        run_configs={
-            "config.yaml": RunConfigSettings(
-                provider_overrides={
-                    "inference": [inference_provider, embedding_provider],
-                    "vector_io": [vector_io_provider],
-                },
-                default_models=[inference_model, embedding_model],
-                default_tool_groups=default_tool_groups,
-            ),
-            "run-with-safety.yaml": RunConfigSettings(
-                provider_overrides={
-                    "inference": [
-                        inference_provider,
-                        embedding_provider,
-                        Provider(
-                            provider_id="meta-reference-safety",
-                            provider_type="inline::meta-reference",
-                            config=MetaReferenceInferenceConfig.sample_run_config(
-                                model="${env.SAFETY_MODEL}",
-                                checkpoint_dir="${env.SAFETY_CHECKPOINT_DIR:=null}",
-                            ),
-                        ),
-                    ],
-                    "vector_io": [vector_io_provider],
-                },
-                default_models=[
-                    inference_model,
-                    safety_model,
-                    embedding_model,
-                ],
-                default_shields=[ShieldInput(shield_id="${env.SAFETY_MODEL}")],
-                default_tool_groups=default_tool_groups,
-            ),
-        },
-        run_config_env_vars={
-            "LLAMA_STACK_PORT": (
-                "8321",
-                "Port for the Llama Stack distribution server",
-            ),
-            "INFERENCE_MODEL": (
-                "meta-llama/Llama-3.2-3B-Instruct",
-                "Inference model loaded into the Meta Reference server",
-            ),
-            "INFERENCE_CHECKPOINT_DIR": (
-                "null",
-                "Directory containing the Meta Reference model checkpoint",
-            ),
-            "SAFETY_MODEL": (
-                "meta-llama/Llama-Guard-3-1B",
-                "Name of the safety (Llama-Guard) model to use",
-            ),
-            "SAFETY_CHECKPOINT_DIR": (
-                "null",
-                "Directory containing the Llama-Guard model checkpoint",
-            ),
-        },
-    )

llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml

@@ -1,155 +0,0 @@
-version: 2
-image_name: meta-reference-gpu
-apis:
-- agents
-- datasetio
-- eval
-- inference
-- safety
-- scoring
-- tool_runtime
-- vector_io
-providers:
-  inference:
-  - provider_id: meta-reference-inference
-    provider_type: inline::meta-reference
-    config:
-      model: ${env.INFERENCE_MODEL}
-      checkpoint_dir: ${env.INFERENCE_CHECKPOINT_DIR:=null}
-      quantization:
-        type: ${env.QUANTIZATION_TYPE:=bf16}
-      model_parallel_size: ${env.MODEL_PARALLEL_SIZE:=0}
-      max_batch_size: ${env.MAX_BATCH_SIZE:=1}
-      max_seq_len: ${env.MAX_SEQ_LEN:=4096}
-  - provider_id: sentence-transformers
-    provider_type: inline::sentence-transformers
-  - provider_id: meta-reference-safety
-    provider_type: inline::meta-reference
-    config:
-      model: ${env.SAFETY_MODEL}
-      checkpoint_dir: ${env.SAFETY_CHECKPOINT_DIR:=null}
-      quantization:
-        type: ${env.QUANTIZATION_TYPE:=bf16}
-      model_parallel_size: ${env.MODEL_PARALLEL_SIZE:=0}
-      max_batch_size: ${env.MAX_BATCH_SIZE:=1}
-      max_seq_len: ${env.MAX_SEQ_LEN:=4096}
-  vector_io:
-  - provider_id: faiss
-    provider_type: inline::faiss
-    config:
-      persistence:
-        namespace: vector_io::faiss
-        backend: kv_default
-  safety:
-  - provider_id: llama-guard
-    provider_type: inline::llama-guard
-    config:
-      excluded_categories: []
-  agents:
-  - provider_id: meta-reference
-    provider_type: inline::meta-reference
-    config:
-      persistence:
-        agent_state:
-          namespace: agents
-          backend: kv_default
-        responses:
-          table_name: responses
-          backend: sql_default
-          max_write_queue_size: 10000
-          num_writers: 4
-  eval:
-  - provider_id: meta-reference
-    provider_type: inline::meta-reference
-    config:
-      kvstore:
-        namespace: eval
-        backend: kv_default
-  datasetio:
-  - provider_id: huggingface
-    provider_type: remote::huggingface
-    config:
-      kvstore:
-        namespace: datasetio::huggingface
-        backend: kv_default
-  - provider_id: localfs
-    provider_type: inline::localfs
-    config:
-      kvstore:
-        namespace: datasetio::localfs
-        backend: kv_default
-  scoring:
-  - provider_id: basic
-    provider_type: inline::basic
-  - provider_id: llm-as-judge
-    provider_type: inline::llm-as-judge
-  - provider_id: braintrust
-    provider_type: inline::braintrust
-    config:
-      openai_api_key: ${env.OPENAI_API_KEY:=}
-  tool_runtime:
-  - provider_id: brave-search
-    provider_type: remote::brave-search
-    config:
-      api_key: ${env.BRAVE_SEARCH_API_KEY:=}
-      max_results: 3
-  - provider_id: tavily-search
-    provider_type: remote::tavily-search
-    config:
-      api_key: ${env.TAVILY_SEARCH_API_KEY:=}
-      max_results: 3
-  - provider_id: rag-runtime
-    provider_type: inline::rag-runtime
-  - provider_id: model-context-protocol
-    provider_type: remote::model-context-protocol
-storage:
-  backends:
-    kv_default:
-      type: kv_sqlite
-      db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db
-    sql_default:
-      type: sql_sqlite
-      db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/sql_store.db
-  stores:
-    metadata:
-      namespace: registry
-      backend: kv_default
-    inference:
-      table_name: inference_store
-      backend: sql_default
-      max_write_queue_size: 10000
-      num_writers: 4
-    conversations:
-      table_name: openai_conversations
-      backend: sql_default
-    prompts:
-      namespace: prompts
-      backend: kv_default
-registered_resources:
-  models:
-  - metadata: {}
-    model_id: ${env.INFERENCE_MODEL}
-    provider_id: meta-reference-inference
-    model_type: llm
-  - metadata: {}
-    model_id: ${env.SAFETY_MODEL}
-    provider_id: meta-reference-safety
-    model_type: llm
-  - metadata:
-      embedding_dimension: 768
-    model_id: nomic-embed-text-v1.5
-    provider_id: sentence-transformers
-    model_type: embedding
-  shields:
-  - shield_id: ${env.SAFETY_MODEL}
-  vector_dbs: []
-  datasets: []
-  scoring_fns: []
-  benchmarks: []
-  tool_groups:
-  - toolgroup_id: builtin::websearch
-    provider_id: tavily-search
-  - toolgroup_id: builtin::rag
-    provider_id: rag-runtime
-server:
-  port: 8321

llama_stack/models/llama/hadamard_utils.py

@@ -1,88 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-import math
-import re
-
-import torch
-from torch import nn
-
-
-def hadamard_transform(x: torch.Tensor) -> torch.Tensor:
-    """Hadamard transform.
-
-    This function performs the Hadamard transform on the input tensor 'x'.
-    The Hadamard transform is a linear transformation that multiplies the input
-    tensor by the Hadamard matrix of dimension n x n, where n is the size of
-    the last dimension of the input tensor.
-    """
-    *_, n = x.shape
-    m = int(math.log2(n))
-    assert n == 1 << m, "n must be a power of 2"
-    x = x[..., None]
-    inv_sqrt2 = 0.5**0.5
-    for _ in range(m):
-        top = x[..., ::2, :] + x[..., 1::2, :]
-        bot = x[..., ::2, :] - x[..., 1::2, :]
-        x = torch.cat((top, bot), dim=-1)
-        x *= inv_sqrt2
-    res = x.squeeze(-2)
-    return res
-
-
-class HadamardModule(torch.nn.Module):
-    """A module that applies the Hadamard transform to the input tensor.
-
-    Args:
-        group_size: The size of the groups that the input tensor will be divided into
-            before applying the Hadamard transform.
-    """
-
-    def __init__(self, group_size: int) -> None:
-        super().__init__()
-        self.group_size = group_size
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        reshape_back = False
-        orig_shape = x.shape
-        if self.group_size != x.shape[-1]:
-            reshape_back = True
-            x = x.reshape(-1, x.shape[-1] // self.group_size, self.group_size)
-        x = hadamard_transform(x)
-        if reshape_back:
-            x = x.reshape(orig_shape)
-        return x
-
-
-def add_hadamard_transform_for_spinquant(model: torch.nn.Module, prefix: str = "") -> None:
-    """
-    Adds a Hadamard transform to the last linear layer of each feedforward network (FFN) in the model.
-    This function recursively traverses the model's children and looks for layers that match the pattern
-    "layers.<digit>.feed_forward.w2", where <digit> is one or more digits. When such a layer is found,
-    it is replaced with a new sequential module that consists of a HadamardModule followed by the original
-    layer. The HadamardModule applies the Hadamard transform to the input tensor.
-
-    See `SpinQuant <https://arxiv.org/abs/2405.16406>_` paper for more details.
-
-    Args:
-        model: An instance of 'torch.nn.Module' (e.g., Transformer model).
-        prefix: A string prefix to add to the full name of each child module.
-
-    Returns:
-        None
-    """
-
-    pattern_last_linear_ffn = r"layers.\d+.feed_forward.w2"
-    for module_name, module in model.named_children():
-        child_full_name = prefix + "." + module_name
-        if re.search(pattern_last_linear_ffn, child_full_name):
-            # Module matching this pattern should be nn.Linear with in_features
-            assert isinstance(module, nn.Linear), f"Expected nn.Linear, got {type(module)}"
-            new_module = nn.Sequential(HadamardModule(group_size=module.in_features), module)
-            del module
-            setattr(model, module_name, new_module)
-        else:
-            add_hadamard_transform_for_spinquant(module, (prefix + "." if prefix else prefix) + module_name)

llama_stack/models/llama/llama3/args.py

@@ -1,74 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from dataclasses import dataclass
-from enum import Enum
-
-
-class QuantizationScheme(Enum):
-    int4_weight_int8_dynamic_activation = "int4_weight_int8_dynamic_activation"
-
-
-@dataclass
-class QuantizationArgs:
-    scheme: QuantizationScheme | None = None
-    group_size: int | None = None
-    spinquant: bool = False
-
-    def __init__(self, **kwargs):
-        for k, v in kwargs.items():
-            if k == "scheme":
-                setattr(self, k, QuantizationScheme(v))
-            else:
-                if hasattr(self, k):
-                    setattr(self, k, v)
-
-
-@dataclass
-class LoRAArgs:
-    rank: int
-    scale: float
-
-
-@dataclass
-class ModelArgs:
-    dim: int = 4096
-    n_layers: int = 32
-    n_heads: int = 32
-    n_kv_heads: int | None = None
-    vocab_size: int = -1
-    multiple_of: int = 256  # make SwiGLU hidden layer size multiple of large power of 2
-    ffn_dim_multiplier: float | None = None
-    norm_eps: float = 1e-5
-    rope_theta: float = 500000
-    use_scaled_rope: bool = False
-
-    max_batch_size: int = 32
-    max_seq_len: int = 2048
-
-    # vision model params
-    vision_chunk_size: int = -1  # image resolution for image models
-    vision_max_num_chunks: int = 4
-    vision_num_cross_attention_layers: int = -1
-
-    quantization_args: QuantizationArgs | None = None
-    lora_args: LoRAArgs | None = None
-
-    def __init__(self, **kwargs):
-        for k, v in kwargs.items():
-            if k == "lora_args":
-                setattr(self, k, LoRAArgs(**v))
-            elif k == "quantization_args":
-                setattr(self, k, QuantizationArgs(**v))
-            else:
-                if hasattr(self, k):
-                    setattr(self, k, v)
-
-        if self.n_kv_heads is None:
-            self.n_kv_heads = self.n_heads
-        assert self.n_kv_heads <= self.n_heads
-        assert self.n_heads % self.n_kv_heads == 0
-        assert self.dim % self.n_heads == 0