pplx-context-mlx 0.1.0__py3-none-any.whl

pplx_context_mlx/__init__.py

@@ -0,0 +1,5 @@
+from .checkpoint import CHECKPOINT_0_6B, REVISION_0_6B, resolve_checkpoint
+from .engine import ContextualEngine, load_model
+from .tokenizer import ContextualTokenizer
+
+__all__ = ["CHECKPOINT_0_6B", "REVISION_0_6B", "ContextualEngine", "ContextualTokenizer", "load_model", "resolve_checkpoint"]

pplx_context_mlx/checkpoint.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import mlx.core as mx
+from huggingface_hub import snapshot_download
+
+from .config import ModelConfig
+from .model import PplxContextualModel
+
+CHECKPOINT_0_6B = "perplexity-ai/pplx-embed-v2-context-0.6b"
+REVISION_0_6B = "6d1fddab0a18478acc0a51edd1e81333fe441095"
+_FILES = ["config.json", "tokenizer.json", "model.safetensors.index.json", "*.safetensors"]
+_DTYPES = {"float32": mx.float32, "float16": mx.float16}
+_KEEP_FP32 = ("contextual_projection.",)
+
+
+def resolve_checkpoint(path_or_repo: str, revision: str | None = None) -> Path:
+    local = Path(path_or_repo).expanduser()
+    if local.is_dir():
+        return local
+    if path_or_repo == CHECKPOINT_0_6B and revision is None:
+        revision = REVISION_0_6B
+    return Path(snapshot_download(path_or_repo, revision=revision, allow_patterns=_FILES))
+
+
+def load_weights(path: Path, dtype: mx.Dtype) -> dict[str, mx.array]:
+    weight_map = json.loads((path / "model.safetensors.index.json").read_text())["weight_map"]
+    weights: dict[str, mx.array] = {}
+    for filename in sorted(set(weight_map.values())):
+        tensors = mx.load(str(path / filename))
+        if not isinstance(tensors, dict):
+            raise TypeError(f"Expected named tensors in {filename}")
+        for key, value in tensors.items():
+            weights[key] = value if key.startswith(_KEEP_FP32) else value.astype(dtype)
+    missing = set(weight_map) - set(weights)
+    if missing:
+        raise ValueError(f"Checkpoint is missing tensors: {sorted(missing)[:5]}")
+    return weights
+
+
+def validate_dtype(dtype: str) -> None:
+    if dtype not in _DTYPES:
+        raise ValueError(f"Unsupported dtype {dtype!r}; choose from {sorted(_DTYPES)}")
+
+
+def load_checkpoint(path: Path, dtype: str = "float32") -> tuple[PplxContextualModel, ModelConfig]:
+    validate_dtype(dtype)
+    config = ModelConfig.from_dict(json.loads((path / "config.json").read_text()))
+    model = PplxContextualModel(config)
+    model.load_weights(list(load_weights(path, _DTYPES[dtype]).items()), strict=True)
+    model.eval()
+    mx.eval(model.parameters())
+    return model, config

pplx_context_mlx/cli.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from .checkpoint import CHECKPOINT_0_6B
+from .engine import load_model
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Embed a JSON list of documents, each a list of chunk strings.")
+    parser.add_argument("input", help="JSON file, or - for stdin")
+    parser.add_argument("--task", choices=["document", "query"], default="document")
+    parser.add_argument("--model", default=CHECKPOINT_0_6B)
+    parser.add_argument("--dtype", default="float32")
+    parser.add_argument("--normalize", action="store_true")
+    args = parser.parse_args()
+    rows = json.loads(sys.stdin.read() if args.input == "-" else Path(args.input).read_text(encoding="utf-8"))
+    engine = load_model(args.model, dtype=args.dtype)
+    encode = engine.encode_queries if args.task == "query" else engine.encode
+    json.dump([row.tolist() for row in encode(rows, normalize_embeddings=args.normalize)], sys.stdout)

pplx_context_mlx/config.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class TextConfig:
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    layer_types: tuple[str, ...]
+    head_dim: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    linear_conv_kernel_dim: int
+    linear_key_head_dim: int
+    linear_num_key_heads: int
+    linear_value_head_dim: int
+    linear_num_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float
+    partial_rotary_factor: float
+    mrope_section: tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class VisionConfig:
+    depth: int
+    hidden_size: int
+    intermediate_size: int
+    in_channels: int
+    num_heads: int
+    num_position_embeddings: int
+    out_hidden_size: int
+    patch_size: int
+    spatial_merge_size: int
+    temporal_patch_size: int
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    text: TextConfig
+    vision: VisionConfig
+    image_token_id: int
+    vision_start_token_id: int
+    vision_end_token_id: int
+    embedding_dim: int
+    query_prefix: str
+    document_prefix: str
+    boundary_marker: str
+    query_length: int
+    document_length: int
+
+    @classmethod
+    def from_dict(cls, raw: dict[str, Any]) -> ModelConfig:
+        if raw.get("model_type") != "pplx_contextual_qwen3_5":
+            raise ValueError(f"Expected a pplx_contextual_qwen3_5 checkpoint, got {raw.get('model_type')!r}")
+        text = raw["text_config"]
+        if text.get("is_causal", True):
+            raise ValueError("Expected a bidirectional checkpoint (text_config.is_causal = false)")
+        kinds = set(text["layer_types"])
+        if not kinds <= {"linear_attention", "full_attention"}:
+            raise ValueError(f"Unsupported layer types: {sorted(kinds)}")
+        if text.get("mlp_only_layers"):
+            raise ValueError("MLP-only layers are not supported")
+        rope = text["rope_parameters"]
+        if rope.get("rope_type") != "default" or not rope.get("mrope_interleaved"):
+            raise ValueError(f"Unsupported rope parameters: {rope}")
+        vision = raw["vision_config"]
+        if vision.get("deepstack_visual_indexes"):
+            raise ValueError("Deepstack visual features are not supported")
+        return cls(
+            vision=VisionConfig(
+                depth=vision["depth"],
+                hidden_size=vision["hidden_size"],
+                intermediate_size=vision["intermediate_size"],
+                in_channels=vision["in_channels"],
+                num_heads=vision["num_heads"],
+                num_position_embeddings=vision["num_position_embeddings"],
+                out_hidden_size=vision["out_hidden_size"],
+                patch_size=vision["patch_size"],
+                spatial_merge_size=vision["spatial_merge_size"],
+                temporal_patch_size=vision["temporal_patch_size"],
+            ),
+            image_token_id=raw["image_token_id"],
+            vision_start_token_id=raw["vision_start_token_id"],
+            vision_end_token_id=raw["vision_end_token_id"],
+            text=TextConfig(
+                vocab_size=text["vocab_size"],
+                hidden_size=text["hidden_size"],
+                intermediate_size=text["intermediate_size"],
+                layer_types=tuple(text["layer_types"]),
+                head_dim=text["head_dim"],
+                num_attention_heads=text["num_attention_heads"],
+                num_key_value_heads=text["num_key_value_heads"],
+                linear_conv_kernel_dim=text["linear_conv_kernel_dim"],
+                linear_key_head_dim=text["linear_key_head_dim"],
+                linear_num_key_heads=text["linear_num_key_heads"],
+                linear_value_head_dim=text["linear_value_head_dim"],
+                linear_num_value_heads=text["linear_num_value_heads"],
+                rms_norm_eps=text["rms_norm_eps"],
+                rope_theta=rope["rope_theta"],
+                partial_rotary_factor=rope["partial_rotary_factor"],
+                mrope_section=tuple(rope["mrope_section"]),
+            ),
+            embedding_dim=raw["embedding_dim"],
+            query_prefix=raw["query_prefix"],
+            document_prefix=raw["document_prefix"],
+            boundary_marker=raw["boundary_marker"],
+            query_length=raw["query_length"],
+            document_length=raw["document_length"],
+        )

pplx_context_mlx/engine.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import mlx.core as mx
+import numpy as np
+from PIL import Image
+
+from .checkpoint import (
+    CHECKPOINT_0_6B,
+    load_checkpoint,
+    resolve_checkpoint,
+    validate_dtype,
+)
+from .config import ModelConfig
+from .images import prepare_images
+from .model import PplxContextualModel, quantize_int8
+from .tokenizer import ContextualTokenizer
+
+
+def _validate(documents: list[list[str]], batch_size: int) -> None:
+    if not isinstance(documents, list) or any(
+        not isinstance(row, list) or any(not isinstance(text, str) for text in row) for row in documents
+    ):
+        raise TypeError("Expected a list of documents, each a list of text chunks")
+    if batch_size < 1:
+        raise ValueError("batch_size must be positive")
+
+
+def _normalize(vectors: mx.array) -> mx.array:
+    return vectors / mx.maximum(mx.linalg.norm(vectors, axis=-1, keepdims=True), 1e-12)
+
+
+class ContextualEngine:
+    """Same call shape as upstream: one (chunk_count, embedding_dim) array per document."""
+
+    def __init__(self, model: PplxContextualModel, tokenizer: ContextualTokenizer, config: ModelConfig):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.config = config
+
+    def project_documents(self, documents: list[list[str]]) -> mx.array:
+        features = self.tokenizer.encode_documents(documents)
+        return self.model.embed_documents(
+            mx.array(features["input_ids"]),
+            mx.array(features["attention_mask"]),
+            mx.array(features["chunk_token_spans"]),
+        )
+
+    def project_queries(self, queries: list[list[str]]) -> mx.array:
+        features = self.tokenizer.encode_queries(queries)
+        return self.model.embed_queries(mx.array(features["input_ids"]), mx.array(features["attention_mask"]))
+
+    def _encode(self, rows: list[list[str]], project, batch_size: int, normalize_embeddings: bool) -> list[np.ndarray]:
+        _validate(rows, batch_size)
+        output: list[np.ndarray] = []
+        for start in range(0, len(rows), batch_size):
+            batch = rows[start : start + batch_size]
+            vectors = quantize_int8(project(batch))
+            if normalize_embeddings:
+                vectors = _normalize(vectors)
+            mx.eval(vectors)
+            matrix = np.array(vectors, dtype=np.float32, copy=True)
+            output.extend(matrix[index, : len(chunks)] for index, chunks in enumerate(batch))
+        return output
+
+    def encode(
+        self, documents: list[list[str]], batch_size: int = 32, *, normalize_embeddings: bool = False
+    ) -> list[np.ndarray]:
+        return self._encode(documents, self.project_documents, batch_size, normalize_embeddings)
+
+    def encode_queries(
+        self, queries: list[list[str]], batch_size: int = 32, *, normalize_embeddings: bool = False
+    ) -> list[np.ndarray]:
+        return self._encode(queries, self.project_queries, batch_size, normalize_embeddings)
+
+    def prepare_images(
+        self, images: list[Image.Image], *, chunks: int = 8, min_pixels: int = 16384, max_pixels: int = 1800000
+    ) -> dict[str, np.ndarray]:
+        return prepare_images(
+            images, self.tokenizer, self.config, chunks=chunks, min_pixels=min_pixels, max_pixels=max_pixels
+        )
+
+    def encode_images(
+        self,
+        images: list[Image.Image],
+        batch_size: int = 1,
+        *,
+        normalize_embeddings: bool = False,
+        chunks: int = 8,
+        min_pixels: int = 16384,
+        max_pixels: int = 1800000,
+    ) -> list[np.ndarray]:
+        """Return one (band_count, embedding_dim) array per image."""
+        if batch_size < 1:
+            raise ValueError("batch_size must be positive")
+        output: list[np.ndarray] = []
+        for start in range(0, len(images), batch_size):
+            features = self.prepare_images(
+                images[start : start + batch_size], chunks=chunks, min_pixels=min_pixels, max_pixels=max_pixels
+            )
+            vectors = quantize_int8(self.model.embed_images(features))
+            if normalize_embeddings:
+                vectors = _normalize(vectors)
+            mx.eval(vectors)
+            matrix = np.array(vectors, dtype=np.float32, copy=True)
+            output.extend(row[mask] for row, mask in zip(matrix, features["chunk_mask"], strict=True))
+        return output
+
+
+def load_model(
+    path_or_repo: str = CHECKPOINT_0_6B,
+    revision: str | None = None,
+    dtype: str = "float32",
+    cache_limit_mb: int = 1024,
+) -> ContextualEngine:
+    validate_dtype(dtype)
+    mx.set_cache_limit(cache_limit_mb * 1024 * 1024)
+    path = resolve_checkpoint(path_or_repo, revision)
+    model, config = load_checkpoint(path, dtype)
+    return ContextualEngine(model, ContextualTokenizer(path, config), config)