nearbits 0.2.0__tar.gz

nearbits-0.2.0/PKG-INFO

@@ -0,0 +1,118 @@
+Metadata-Version: 2.4
+Name: nearbits
+Version: 0.2.0
+Summary: Hybrid low-bit compression for Hugging Face models on low-RAM devices.
+Author: Codex
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: transformers>=4.40.0
+Requires-Dist: accelerate>=0.28.0
+Requires-Dist: sentencepiece>=0.1.99
+Requires-Dist: psutil>=5.9.0
+
+# NearBit
+
+NearBit is a Python package for compressing Hugging Face models with a single strategy: `hybrid_lowbit`.
+
+The package keeps high-impact parts in `int8`, compresses the bulk of linear layers to `binary` or `ternary`, rescues outlier weights in `int8`, and targets the lowest practical RAM footprint without retraining.
+
+## What NearBit does
+
+- Compresses `nn.Linear` layers with `binary` or `ternary` storage plus group-wise scales
+- Keeps sensitive linear layers in `int8`
+- Compresses `nn.Embedding` tables to `int8`
+- Works directly on existing Hugging Face models
+- Exposes simple chat and prompt benchmarking helpers
+
+## Install
+
+Install PyTorch first for your platform, then install NearBit.
+
+```bash
+pip install torch
+pip install nearbit
+```
+
+For local development from this repo:
+
+```bash
+pip install torch
+pip install -e .
+```
+
+## Quick Start
+
+```python
+from nearbit import HybridLowBitConfig, chat, compress_hf_model, load_hf_model, model_report
+
+model, tokenizer, model_kind = load_hf_model("Qwen/Qwen2.5-1.5B-Instruct")
+print(model_report(model))
+
+config = HybridLowBitConfig(
+    group_size=128,
+    binary_outlier_fraction=0.002,
+    ternary_outlier_fraction=0.001,
+    binary_nmse_threshold=0.045,
+    ternary_nmse_threshold=0.02,
+)
+
+compressed_model = compress_hf_model(model, config)
+print(model_report(compressed_model))
+
+reply = chat(compressed_model, tokenizer, "Who are you?")
+print(reply)
+```
+
+## Google Colab Usage
+
+Use `examples/google_collab_hybrid_lowbit.py` as the runnable Colab example.
+
+Typical Colab cells:
+
+```python
+!pip install torch
+!pip install nearbit
+```
+
+```python
+from nearbit import HybridLowBitConfig, benchmark_prompts, compress_hf_model, load_hf_model, model_report
+
+MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
+PROMPTS = [
+    "hi",
+    "hello",
+    "who are you?",
+    "what is 2 + 2?",
+    "write one sentence about India.",
+    "explain model compression in simple words.",
+]
+
+model, tokenizer, model_kind = load_hf_model(MODEL_ID)
+print("Before:", model_report(model))
+
+compressed_model = compress_hf_model(model, HybridLowBitConfig())
+print("After:", model_report(compressed_model))
+
+rows = benchmark_prompts(compressed_model, tokenizer, PROMPTS)
+for row in rows:
+    print(row["prompt"])
+    print(row["response"])
+    print(row["latency_sec"])
+```
+
+## Public API
+
+- `load_hf_model(model_id, trust_remote_code=False)`
+- `compress_hf_model(model, config=None)`
+- `model_size_mb(model)`
+- `model_report(model)`
+- `chat(model, tokenizer, prompt, ...)`
+- `benchmark_prompts(model, tokenizer, prompts, ...)`
+- `HybridLowBitConfig(...)`
+
+## Notes
+
+- NearBit is centered on one executable path: `hybrid_lowbit`
+- There is no public `int8` fallback API in the package surface
+- Quality depends on the model architecture and compression thresholds
+- The strongest RAM reduction comes from the hybrid path, not from pure `int8`

nearbits-0.2.0/README.md

@@ -0,0 +1,106 @@
+# NearBit
+
+NearBit is a Python package for compressing Hugging Face models with a single strategy: `hybrid_lowbit`.
+
+The package keeps high-impact parts in `int8`, compresses the bulk of linear layers to `binary` or `ternary`, rescues outlier weights in `int8`, and targets the lowest practical RAM footprint without retraining.
+
+## What NearBit does
+
+- Compresses `nn.Linear` layers with `binary` or `ternary` storage plus group-wise scales
+- Keeps sensitive linear layers in `int8`
+- Compresses `nn.Embedding` tables to `int8`
+- Works directly on existing Hugging Face models
+- Exposes simple chat and prompt benchmarking helpers
+
+## Install
+
+Install PyTorch first for your platform, then install NearBit.
+
+```bash
+pip install torch
+pip install nearbit
+```
+
+For local development from this repo:
+
+```bash
+pip install torch
+pip install -e .
+```
+
+## Quick Start
+
+```python
+from nearbit import HybridLowBitConfig, chat, compress_hf_model, load_hf_model, model_report
+
+model, tokenizer, model_kind = load_hf_model("Qwen/Qwen2.5-1.5B-Instruct")
+print(model_report(model))
+
+config = HybridLowBitConfig(
+    group_size=128,
+    binary_outlier_fraction=0.002,
+    ternary_outlier_fraction=0.001,
+    binary_nmse_threshold=0.045,
+    ternary_nmse_threshold=0.02,
+)
+
+compressed_model = compress_hf_model(model, config)
+print(model_report(compressed_model))
+
+reply = chat(compressed_model, tokenizer, "Who are you?")
+print(reply)
+```
+
+## Google Colab Usage
+
+Use `examples/google_collab_hybrid_lowbit.py` as the runnable Colab example.
+
+Typical Colab cells:
+
+```python
+!pip install torch
+!pip install nearbit
+```
+
+```python
+from nearbit import HybridLowBitConfig, benchmark_prompts, compress_hf_model, load_hf_model, model_report
+
+MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
+PROMPTS = [
+    "hi",
+    "hello",
+    "who are you?",
+    "what is 2 + 2?",
+    "write one sentence about India.",
+    "explain model compression in simple words.",
+]
+
+model, tokenizer, model_kind = load_hf_model(MODEL_ID)
+print("Before:", model_report(model))
+
+compressed_model = compress_hf_model(model, HybridLowBitConfig())
+print("After:", model_report(compressed_model))
+
+rows = benchmark_prompts(compressed_model, tokenizer, PROMPTS)
+for row in rows:
+    print(row["prompt"])
+    print(row["response"])
+    print(row["latency_sec"])
+```
+
+## Public API
+
+- `load_hf_model(model_id, trust_remote_code=False)`
+- `compress_hf_model(model, config=None)`
+- `model_size_mb(model)`
+- `model_report(model)`
+- `chat(model, tokenizer, prompt, ...)`
+- `benchmark_prompts(model, tokenizer, prompts, ...)`
+- `HybridLowBitConfig(...)`
+
+## Notes
+
+- NearBit is centered on one executable path: `hybrid_lowbit`
+- There is no public `int8` fallback API in the package surface
+- Quality depends on the model architecture and compression thresholds
+- The strongest RAM reduction comes from the hybrid path, not from pure `int8`

nearbits-0.2.0/pyproject.toml

@@ -0,0 +1,25 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "nearbits"
+version = "0.2.0"
+description = "Hybrid low-bit compression for Hugging Face models on low-RAM devices."
+readme = "README.md"
+requires-python = ">=3.10"
+authors = [
+  { name = "Codex" }
+]
+dependencies = [
+  "transformers>=4.40.0",
+  "accelerate>=0.28.0",
+  "sentencepiece>=0.1.99",
+  "psutil>=5.9.0"
+]
+
+[tool.setuptools]
+package-dir = { "" = "src" }
+
+[tool.setuptools.packages.find]
+where = ["src"]

nearbits-0.2.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

nearbits-0.2.0/src/nearbit/__init__.py

@@ -0,0 +1,14 @@
+from .api import benchmark_prompts, chat, compress_hf_model, load_hf_model, model_report, model_size_mb
+from .integration import HybridLowBitConfig, detect_model_dtypes, estimate_effective_bits
+
+__all__ = [
+    "HybridLowBitConfig",
+    "benchmark_prompts",
+    "chat",
+    "compress_hf_model",
+    "detect_model_dtypes",
+    "estimate_effective_bits",
+    "load_hf_model",
+    "model_report",
+    "model_size_mb",
+]

nearbits-0.2.0/src/nearbit/analysis/__init__.py

@@ -0,0 +1,4 @@
+from .analyzer import Analyzer
+from .introspection import ModelInspector
+
+__all__ = ["Analyzer", "ModelInspector"]

nearbits-0.2.0/src/nearbit/analysis/analyzer.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from statistics import mean
+from typing import Any
+
+from ..types import CalibrationSample, HardwareProfile, LayerSensitivity, PrecisionKind, SensitivityReport
+from .introspection import ModelInspector
+
+
+class Analyzer:
+    def __init__(self, hardware_profile: HardwareProfile) -> None:
+        self.hardware_profile = hardware_profile
+        self.inspector = ModelInspector()
+
+    def analyze(self, model: Any, calibration_set: Iterable[CalibrationSample]) -> SensitivityReport:
+        samples = list(calibration_set)
+        model_name, model_type = self.inspector.identify_model(model)
+        layers = self.inspector.inspect_layers(model)
+        sample_factor = max(len(samples), 1)
+
+        sensitivities: list[LayerSensitivity] = []
+        for index, layer in enumerate(layers):
+            size_factor = min(layer.parameter_count / 1_000_000, 4.0)
+            position_penalty = 0.8 if index in (0, len(layers) - 1) else 0.0
+            fragile_penalty = 1.5 if layer.is_fragile else 0.0
+            binary_bonus = -0.3 if layer.supports_binary else 0.8
+            score = round(0.4 + size_factor + position_penalty + fragile_penalty + binary_bonus, 4)
+
+            min_act = round(-1.0 - 0.02 * index, 4)
+            max_act = round(1.0 + 0.03 * index, 4)
+            outlier_fraction = round(min(0.01, 0.001 + size_factor / 100 + index / (10_000 * sample_factor)), 6)
+            estimated_binary_error = round(score * 0.018 + outlier_fraction * 10, 6)
+
+            recommended = self._recommended_precision(layer.is_fragile, layer.supports_binary, estimated_binary_error)
+            sensitivities.append(
+                LayerSensitivity(
+                    name=layer.name,
+                    op_type=layer.op_type,
+                    sensitivity_score=score,
+                    activation_range=(min_act, max_act),
+                    outlier_fraction=outlier_fraction,
+                    estimated_binary_error=estimated_binary_error,
+                    is_fragile=layer.is_fragile,
+                    supports_binary=layer.supports_binary,
+                    recommended_precision=recommended,
+                )
+            )
+
+        baseline_metric = round(max(0.0, 1.0 - mean(layer.estimated_binary_error for layer in sensitivities) / 2), 6)
+        return SensitivityReport(
+            model_name=model_name,
+            model_type=model_type,
+            layers=sensitivities,
+            baseline_metric=baseline_metric,
+            metadata={
+                "calibration_samples": len(samples),
+                "hardware_target": self.hardware_profile.device_name,
+                "accelerator": self.hardware_profile.accelerator,
+            },
+        )
+
+    def _recommended_precision(
+        self,
+        is_fragile: bool,
+        supports_binary: bool,
+        estimated_binary_error: float,
+    ) -> PrecisionKind:
+        if is_fragile:
+            return PrecisionKind.INT8
+        if not supports_binary:
+            return PrecisionKind.INT4 if estimated_binary_error < 0.08 else PrecisionKind.INT8
+        if estimated_binary_error < 0.03:
+            return PrecisionKind.BINARY
+        if estimated_binary_error < 0.06:
+            return PrecisionKind.BINARY_OUTLIERS
+        if estimated_binary_error < 0.09:
+            return PrecisionKind.BINARY_LOW_RANK
+        return PrecisionKind.INT4

nearbits-0.2.0/src/nearbit/analysis/introspection.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import Any, Iterable
+
+from ..types import LayerProfile
+
+
+class ModelInspector:
+    """Adapter-style model inspection with safe fallbacks for existing models."""
+
+    FRAGILE_OPS = {"embedding", "norm", "layernorm", "rmsnorm", "head", "lm_head"}
+    SAFE_BINARY_OPS = {"linear", "conv1d", "conv2d", "matmul"}
+
+    def inspect_layers(self, model: Any) -> list[LayerProfile]:
+        if hasattr(model, "nearbit_layers"):
+            return list(model.nearbit_layers())
+        if hasattr(model, "named_modules"):
+            return self._from_named_modules(model.named_modules())
+        if isinstance(model, dict) and "layers" in model:
+            return self._from_mapping(model["layers"])
+        raise TypeError("Unsupported model format. Provide nearbit_layers(), named_modules(), or a layer mapping.")
+
+    def identify_model(self, model: Any) -> tuple[str, str]:
+        model_name = getattr(model, "name", model.__class__.__name__)
+        model_type = getattr(model, "model_type", "transformer")
+        return model_name, model_type
+
+    def _from_named_modules(self, modules: Iterable[tuple[str, Any]]) -> list[LayerProfile]:
+        layers: list[LayerProfile] = []
+        for name, module in modules:
+            if not name:
+                continue
+            op_type = module.__class__.__name__.lower()
+            parameter_count = int(getattr(module, "parameter_count", 0) or self._safe_numel(module))
+            fragile = self._is_fragile(name, op_type)
+            supports_binary = self._supports_binary(op_type) and not fragile
+            layers.append(
+                LayerProfile(
+                    name=name,
+                    op_type=op_type,
+                    parameter_count=parameter_count,
+                    is_fragile=fragile,
+                    supports_binary=supports_binary,
+                )
+            )
+        return layers
+
+    def _from_mapping(self, layers_map: Iterable[dict[str, Any]]) -> list[LayerProfile]:
+        layers: list[LayerProfile] = []
+        for item in layers_map:
+            name = item["name"]
+            op_type = str(item.get("op_type", "linear")).lower()
+            fragile = bool(item.get("is_fragile", self._is_fragile(name, op_type)))
+            supports_binary = bool(item.get("supports_binary", self._supports_binary(op_type) and not fragile))
+            layers.append(
+                LayerProfile(
+                    name=name,
+                    op_type=op_type,
+                    parameter_count=int(item.get("parameter_count", 0)),
+                    output_shape=tuple(item.get("output_shape", ())),
+                    is_fragile=fragile,
+                    supports_binary=supports_binary,
+                )
+            )
+        return layers
+
+    def _is_fragile(self, name: str, op_type: str) -> bool:
+        name_lower = name.lower()
+        return any(token in name_lower or token in op_type for token in self.FRAGILE_OPS)
+
+    def _supports_binary(self, op_type: str) -> bool:
+        return any(token in op_type for token in self.SAFE_BINARY_OPS)
+
+    def _safe_numel(self, module: Any) -> int:
+        parameters = getattr(module, "parameters", None)
+        if parameters is None:
+            return 0
+        total = 0
+        for param in parameters():
+            total += int(getattr(param, "numel", lambda: 0)())
+        return total

nearbits-0.2.0/src/nearbit/api.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import torch
+
+from .integration import (
+    HybridLowBitConfig,
+    convert_model_to_hybrid_lowbit,
+    detect_model_dtypes,
+    estimate_effective_bits,
+)
+
+
+def load_hf_model(model_id: str, trust_remote_code: bool = False):
+    from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=trust_remote_code)
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    errors: list[str] = []
+    for model_kind, loader in (("causal", AutoModelForCausalLM), ("seq2seq", AutoModelForSeq2SeqLM)):
+        try:
+            model = loader.from_pretrained(
+                model_id,
+                trust_remote_code=trust_remote_code,
+                low_cpu_mem_usage=True,
+            )
+            model.eval()
+            return model, tokenizer, model_kind
+        except Exception as exc:
+            errors.append(f"{model_kind}: {exc}")
+    raise RuntimeError(f"Could not load model '{model_id}'. Details: {' | '.join(errors)}")
+
+
+def compress_hf_model(model: Any, config: HybridLowBitConfig | None = None):
+    return convert_model_to_hybrid_lowbit(model, config or HybridLowBitConfig())
+
+
+def model_size_mb(model: Any) -> float:
+    total_bytes = 0
+    for tensor in model.state_dict().values():
+        if hasattr(tensor, "element_size") and hasattr(tensor, "nelement"):
+            total_bytes += tensor.element_size() * tensor.nelement()
+    return round(total_bytes / (1024 ** 2), 4)
+
+
+def model_report(model: Any) -> dict[str, Any]:
+    return {
+        "model_size_mb": model_size_mb(model),
+        "effective_bits": estimate_effective_bits(model),
+        "dtypes": detect_model_dtypes(model),
+        "lowbit_summary": dict(getattr(model, "nearbit_lowbit_summary", {})),
+    }
+
+
+def chat(
+    model: Any,
+    tokenizer: Any,
+    prompt: str,
+    *,
+    system_prompt: str | None = "You are a concise, helpful assistant.",
+    max_new_tokens: int = 128,
+    do_sample: bool = False,
+    temperature: float = 0.7,
+    top_p: float = 0.95,
+) -> str:
+    from transformers import GenerationConfig
+
+    inputs = _build_inputs(tokenizer, prompt, system_prompt)
+    generation_kwargs = {
+        "max_new_tokens": max_new_tokens,
+        "do_sample": do_sample,
+        "pad_token_id": tokenizer.pad_token_id,
+        "eos_token_id": tokenizer.eos_token_id,
+    }
+    if do_sample:
+        generation_kwargs["temperature"] = temperature
+        generation_kwargs["top_p"] = top_p
+    generation_config = GenerationConfig(**generation_kwargs)
+
+    with torch.inference_mode():
+        output = model.generate(**inputs, generation_config=generation_config)
+
+    is_encoder_decoder = bool(getattr(getattr(model, "config", None), "is_encoder_decoder", False))
+    if is_encoder_decoder:
+        generated_ids = output[0]
+    else:
+        input_len = inputs["input_ids"].shape[-1]
+        generated_ids = output[0][input_len:]
+    return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+
+
+def benchmark_prompts(
+    model: Any,
+    tokenizer: Any,
+    prompts: list[str],
+    *,
+    system_prompt: str | None = "You are a concise, helpful assistant.",
+    max_new_tokens: int = 64,
+    do_sample: bool = False,
+    temperature: float = 0.7,
+    top_p: float = 0.95,
+) -> list[dict[str, Any]]:
+    rows: list[dict[str, Any]] = []
+    for prompt in prompts:
+        started = time.perf_counter()
+        response = chat(
+            model,
+            tokenizer,
+            prompt,
+            system_prompt=system_prompt,
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_p=top_p,
+        )
+        elapsed = time.perf_counter() - started
+        token_count = len(tokenizer.encode(response, add_special_tokens=False))
+        rows.append(
+            {
+                "prompt": prompt,
+                "response": response,
+                "latency_sec": round(elapsed, 4),
+                "generated_tokens": token_count,
+                "tokens_per_sec": round(token_count / elapsed, 4) if elapsed > 0 else None,
+            }
+        )
+    return rows
+
+
+def _build_inputs(tokenizer: Any, prompt: str, system_prompt: str | None = None):
+    if getattr(tokenizer, "chat_template", None):
+        messages = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+        rendered = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        return tokenizer(rendered, return_tensors="pt")
+    return tokenizer(prompt, return_tensors="pt")

nearbits-0.2.0/src/nearbit/compression/__init__.py

@@ -0,0 +1,3 @@
+from .compressor import Compressor
+
+__all__ = ["Compressor"]

nearbits-0.2.0/src/nearbit/compression/activations.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from ..types import PrecisionKind
+
+
+def activation_policy(model_type: str) -> dict[str, str]:
+    if model_type == "transformer":
+        return {"default_precision": PrecisionKind.INT8.value, "fallback_precision": PrecisionKind.INT4.value}
+    return {"default_precision": PrecisionKind.INT8.value, "fallback_precision": PrecisionKind.INT8.value}

nearbits-0.2.0/src/nearbit/compression/binary.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import Iterable
+
+
+def pack_signs(weights: Iterable[float], group_size: int) -> dict[str, object]:
+    values = list(weights)
+    packed = [1 if value >= 0 else 0 for value in values]
+    scales = []
+    for index in range(0, len(values), group_size):
+        group = values[index : index + group_size]
+        mean_abs = sum(abs(value) for value in group) / max(len(group), 1)
+        scales.append(mean_abs)
+    return {
+        "bitpack": packed,
+        "group_size": group_size,
+        "scales": scales,
+    }

nearbits-0.2.0/src/nearbit/compression/compressor.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import Any
+
+from ..packaging.builder import ArtifactBuilder
+from ..types import CompressionArtifact, CompressionPlan, PackedLayer, PrecisionKind
+from .binary import pack_signs
+from .kv_cache import kv_cache_policy
+from .outliers import extract_outliers
+from .residual import build_low_rank_residual
+
+
+class Compressor:
+    def compress(self, model: Any, compression_plan: CompressionPlan) -> CompressionArtifact:
+        weights_by_layer = self._read_weights(model)
+        packed_layers: list[PackedLayer] = []
+
+        for layer_plan in compression_plan.layers:
+            weights = weights_by_layer.get(layer_plan.name, [0.0])
+            payload = self._compress_layer(weights, layer_plan.precision, layer_plan.group_size, layer_plan.outlier_fraction, layer_plan.residual_rank)
+            packed_layers.append(
+                PackedLayer(
+                    name=layer_plan.name,
+                    precision=layer_plan.precision,
+                    payload=payload,
+                )
+            )
+
+        return ArtifactBuilder().build(
+            compression_plan=compression_plan,
+            packed_layers=packed_layers,
+            kv_policy=kv_cache_policy(compression_plan.model_type),
+        )
+
+    def _compress_layer(
+        self,
+        weights: list[float],
+        precision: PrecisionKind,
+        group_size: int,
+        outlier_fraction: float,
+        residual_rank: int,
+    ) -> dict[str, object]:
+        if precision == PrecisionKind.BINARY:
+            return pack_signs(weights, group_size)
+        if precision == PrecisionKind.BINARY_OUTLIERS:
+            return {
+                "binary": pack_signs(weights, group_size),
+                "outliers": extract_outliers(weights, outlier_fraction),
+            }
+        if precision == PrecisionKind.BINARY_LOW_RANK:
+            return {
+                "binary": pack_signs(weights, group_size),
+                "outliers": extract_outliers(weights, outlier_fraction),
+                "residual": build_low_rank_residual(weights, residual_rank),
+            }
+        return {
+            "values": weights,
+            "stored_precision": precision.value,
+        }
+
+    def _read_weights(self, model: Any) -> dict[str, list[float]]:
+        if hasattr(model, "nearbit_weights"):
+            return dict(model.nearbit_weights())
+        if isinstance(model, dict) and "weights" in model:
+            return {str(key): list(value) for key, value in model["weights"].items()}
+        return {}