nsos-mamba 0.1.0__tar.gz

nsos_mamba-0.1.0/PKG-INFO

@@ -0,0 +1,228 @@
+Metadata-Version: 2.4
+Name: nsos-mamba
+Version: 0.1.0
+Summary: Wrapper Python isolado para o módulo Mamba do NSOS com streaming incremental padrão.
+Author: Oxta
+Project-URL: Homepage, https://github.com/vitorGgC569/reimagined
+Project-URL: Documentation, https://github.com/vitorGgC569/reimagined/tree/nsos-gpu-phases12/OXN/nsos/docs
+Project-URL: Source, https://github.com/vitorGgC569/reimagined
+Keywords: mamba,ssm,nsos,edge-ai,streaming
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.23
+
+# NSOS Mamba
+
+`nsos-mamba` é um wrapper Python simples para usar apenas o módulo Mamba do
+NSOS. Ele configura automaticamente o modelo em modo Mamba-2 faithful, ativa
+streaming incremental por padrão e oferece utilitários pequenos para treinar
+pares `prompt -> resposta`.
+
+Este pacote não instala o runtime nativo `nsos_ext`. Antes de usar, compile ou
+copie o `nsos_ext.so`/`nsos_ext.pyd` e deixe-o acessível via `PYTHONPATH` ou
+pela variável `NSOS_EXT_PATH`.
+
+## Instalação
+
+```bash
+pip install nsos-mamba
+```
+
+Se o `nsos_ext` estiver em uma pasta separada:
+
+```bash
+export NSOS_EXT_PATH=/caminho/para/nsos_ext
+```
+
+No Colab:
+
+```python
+import sys, os
+
+sys.path.insert(0, "/content/nsos_ext_mambavs")
+# ou:
+os.environ["NSOS_EXT_PATH"] = "/content/nsos_ext_mambavs"
+```
+
+## Configurações principais
+
+```python
+from nsos_mamba import MambaModuleConfig
+
+cfg = MambaModuleConfig(
+    vocab_size=128,
+    num_layers=12,
+    d_model=128,
+    device="auto",
+    streaming=True,
+)
+```
+
+Campos mais usados:
+
+- `vocab_size`: tamanho do vocabulário.
+- `num_layers`: número de camadas Mamba.
+- `d_model`: largura do modelo.
+- `d_state`: tamanho do estado SSM. Padrão: `64`.
+- `max_context_tokens`: limite de contexto. Padrão: `4096`.
+- `device`: `"auto"`, `"gpu"`, `"cuda"` ou `"cpu"`.
+- `streaming`: ativa decode incremental. Padrão: `True`.
+- `allow_cpu_fallback`: se a GPU falhar, cai para CPU. Padrão: `True`.
+
+Variáveis de ambiente equivalentes:
+
+```bash
+NSOS_MAMBA_DEVICE=auto
+NSOS_MAMBA_STREAMING=1
+NSOS_MAMBA_LAYERS=12
+NSOS_MAMBA_DMODEL=128
+NSOS_MAMBA_DSTATE=64
+NSOS_MAMBA_CONTEXT=4096
+```
+
+## Exemplo mínimo: "Quem é você?" -> "Oxta"
+
+```python
+from nsos_mamba import CharTokenizer, MambaModuleConfig, NSOSMamba, TextPair
+
+pairs = [
+    TextPair("Quem é você?", "Oxta"),
+    TextPair("quem é você?", "Oxta"),
+    TextPair("Quem e voce?", "Oxta"),
+]
+
+tok = CharTokenizer.from_texts(
+    [p.prompt for p in pairs] + [p.answer for p in pairs]
+)
+
+cfg = MambaModuleConfig(
+    vocab_size=tok.vocab_size,
+    num_layers=12,
+    d_model=128,
+    device="auto",
+    streaming=True,
+)
+
+mamba = NSOSMamba(cfg)
+
+mamba.fit_text_pairs(
+    pairs,
+    tok,
+    steps=1200,
+    batch_size=16,
+    learning_rate=3e-3,
+    callback=lambda step, loss: print(f"[train] step {step} loss~{loss:.4f}"),
+)
+
+result = mamba.generate_text("Quem é você?", tok, max_new_tokens=8)
+
+print("resposta:", repr(result.text))
+print("tokens/s:", result.decode_tokens_per_sec)
+```
+
+Saída esperada:
+
+```text
+resposta: 'Oxta'
+```
+
+## Streaming incremental
+
+O streaming incremental é o padrão. O fluxo usado é:
+
+1. processa o prompt uma vez (`prefill`);
+2. gera cada token novo passando apenas o último token;
+3. reaproveita o estado interno do modelo.
+
+Benchmark:
+
+```python
+prompt_ids = tok.encode("Quem é você?", bos=True)
+bench = mamba.benchmark_decode(prompt_ids, max_new_tokens=1024, repeats=5)
+bench.print()
+```
+
+Métricas:
+
+- `prefill_ms`: tempo para processar o prompt.
+- `first_incremental_ms`: primeiro passo incremental após o prefill.
+- `decode_tokens_per_sec`: tokens/s do decode incremental.
+- `end_to_end_tokens_per_sec`: tokens/s incluindo prefill.
+
+## Datasets
+
+### Lista em memória
+
+```python
+from nsos_mamba import TextPair
+
+pairs = [
+    TextPair(prompt="Quem é você?", answer="Oxta"),
+    TextPair(prompt="Qual seu nome?", answer="Oxta"),
+]
+```
+
+### JSON
+
+Arquivo `dataset.json`:
+
+```json
+[
+  {"prompt": "Quem é você?", "answer": "Oxta"},
+  {"prompt": "Qual seu nome?", "answer": "Oxta"}
+]
+```
+
+Uso:
+
+```python
+from nsos_mamba import load_text_pairs_json
+
+pairs = load_text_pairs_json("dataset.json")
+```
+
+Também aceita:
+
+```json
+{"data": [{"prompt": "...", "answer": "..."}]}
+```
+
+ou:
+
+```json
+{"examples": [{"prompt": "...", "answer": "..."}]}
+```
+
+### JSONL
+
+Arquivo `dataset.jsonl`:
+
+```jsonl
+{"prompt": "Quem é você?", "answer": "Oxta"}
+{"prompt": "Qual seu nome?", "answer": "Oxta"}
+```
+
+Uso:
+
+```python
+from nsos_mamba import load_text_pairs_jsonl
+
+pairs = load_text_pairs_jsonl("dataset.jsonl")
+```
+
+## Observações
+
+- `CharTokenizer` é para testes, demos e dados sintéticos pequenos.
+- Para treino real, use um tokenizer próprio do projeto.
+- O pacote expõe o Mamba isolado; não ativa Attention, MoE, KAN, TTT, CHRASS,
+  MCTS ou memória externa.
+- O pacote é uma camada de uso sobre `nsos_ext`, não uma reimplementação em
+  Python.

nsos_mamba-0.1.0/README.md

@@ -0,0 +1,207 @@
+# NSOS Mamba
+
+`nsos-mamba` é um wrapper Python simples para usar apenas o módulo Mamba do
+NSOS. Ele configura automaticamente o modelo em modo Mamba-2 faithful, ativa
+streaming incremental por padrão e oferece utilitários pequenos para treinar
+pares `prompt -> resposta`.
+
+Este pacote não instala o runtime nativo `nsos_ext`. Antes de usar, compile ou
+copie o `nsos_ext.so`/`nsos_ext.pyd` e deixe-o acessível via `PYTHONPATH` ou
+pela variável `NSOS_EXT_PATH`.
+
+## Instalação
+
+```bash
+pip install nsos-mamba
+```
+
+Se o `nsos_ext` estiver em uma pasta separada:
+
+```bash
+export NSOS_EXT_PATH=/caminho/para/nsos_ext
+```
+
+No Colab:
+
+```python
+import sys, os
+
+sys.path.insert(0, "/content/nsos_ext_mambavs")
+# ou:
+os.environ["NSOS_EXT_PATH"] = "/content/nsos_ext_mambavs"
+```
+
+## Configurações principais
+
+```python
+from nsos_mamba import MambaModuleConfig
+
+cfg = MambaModuleConfig(
+    vocab_size=128,
+    num_layers=12,
+    d_model=128,
+    device="auto",
+    streaming=True,
+)
+```
+
+Campos mais usados:
+
+- `vocab_size`: tamanho do vocabulário.
+- `num_layers`: número de camadas Mamba.
+- `d_model`: largura do modelo.
+- `d_state`: tamanho do estado SSM. Padrão: `64`.
+- `max_context_tokens`: limite de contexto. Padrão: `4096`.
+- `device`: `"auto"`, `"gpu"`, `"cuda"` ou `"cpu"`.
+- `streaming`: ativa decode incremental. Padrão: `True`.
+- `allow_cpu_fallback`: se a GPU falhar, cai para CPU. Padrão: `True`.
+
+Variáveis de ambiente equivalentes:
+
+```bash
+NSOS_MAMBA_DEVICE=auto
+NSOS_MAMBA_STREAMING=1
+NSOS_MAMBA_LAYERS=12
+NSOS_MAMBA_DMODEL=128
+NSOS_MAMBA_DSTATE=64
+NSOS_MAMBA_CONTEXT=4096
+```
+
+## Exemplo mínimo: "Quem é você?" -> "Oxta"
+
+```python
+from nsos_mamba import CharTokenizer, MambaModuleConfig, NSOSMamba, TextPair
+
+pairs = [
+    TextPair("Quem é você?", "Oxta"),
+    TextPair("quem é você?", "Oxta"),
+    TextPair("Quem e voce?", "Oxta"),
+]
+
+tok = CharTokenizer.from_texts(
+    [p.prompt for p in pairs] + [p.answer for p in pairs]
+)
+
+cfg = MambaModuleConfig(
+    vocab_size=tok.vocab_size,
+    num_layers=12,
+    d_model=128,
+    device="auto",
+    streaming=True,
+)
+
+mamba = NSOSMamba(cfg)
+
+mamba.fit_text_pairs(
+    pairs,
+    tok,
+    steps=1200,
+    batch_size=16,
+    learning_rate=3e-3,
+    callback=lambda step, loss: print(f"[train] step {step} loss~{loss:.4f}"),
+)
+
+result = mamba.generate_text("Quem é você?", tok, max_new_tokens=8)
+
+print("resposta:", repr(result.text))
+print("tokens/s:", result.decode_tokens_per_sec)
+```
+
+Saída esperada:
+
+```text
+resposta: 'Oxta'
+```
+
+## Streaming incremental
+
+O streaming incremental é o padrão. O fluxo usado é:
+
+1. processa o prompt uma vez (`prefill`);
+2. gera cada token novo passando apenas o último token;
+3. reaproveita o estado interno do modelo.
+
+Benchmark:
+
+```python
+prompt_ids = tok.encode("Quem é você?", bos=True)
+bench = mamba.benchmark_decode(prompt_ids, max_new_tokens=1024, repeats=5)
+bench.print()
+```
+
+Métricas:
+
+- `prefill_ms`: tempo para processar o prompt.
+- `first_incremental_ms`: primeiro passo incremental após o prefill.
+- `decode_tokens_per_sec`: tokens/s do decode incremental.
+- `end_to_end_tokens_per_sec`: tokens/s incluindo prefill.
+
+## Datasets
+
+### Lista em memória
+
+```python
+from nsos_mamba import TextPair
+
+pairs = [
+    TextPair(prompt="Quem é você?", answer="Oxta"),
+    TextPair(prompt="Qual seu nome?", answer="Oxta"),
+]
+```
+
+### JSON
+
+Arquivo `dataset.json`:
+
+```json
+[
+  {"prompt": "Quem é você?", "answer": "Oxta"},
+  {"prompt": "Qual seu nome?", "answer": "Oxta"}
+]
+```
+
+Uso:
+
+```python
+from nsos_mamba import load_text_pairs_json
+
+pairs = load_text_pairs_json("dataset.json")
+```
+
+Também aceita:
+
+```json
+{"data": [{"prompt": "...", "answer": "..."}]}
+```
+
+ou:
+
+```json
+{"examples": [{"prompt": "...", "answer": "..."}]}
+```
+
+### JSONL
+
+Arquivo `dataset.jsonl`:
+
+```jsonl
+{"prompt": "Quem é você?", "answer": "Oxta"}
+{"prompt": "Qual seu nome?", "answer": "Oxta"}
+```
+
+Uso:
+
+```python
+from nsos_mamba import load_text_pairs_jsonl
+
+pairs = load_text_pairs_jsonl("dataset.jsonl")
+```
+
+## Observações
+
+- `CharTokenizer` é para testes, demos e dados sintéticos pequenos.
+- Para treino real, use um tokenizer próprio do projeto.
+- O pacote expõe o Mamba isolado; não ativa Attention, MoE, KAN, TTT, CHRASS,
+  MCTS ou memória externa.
+- O pacote é uma camada de uso sobre `nsos_ext`, não uma reimplementação em
+  Python.

nsos_mamba-0.1.0/nsos_mamba/__init__.py

@@ -0,0 +1,37 @@
+"""Standalone NSOS Mamba Python wrapper.
+
+This package is intentionally thin: it does not implement a new architecture.
+It builds a pure faithful-Mamba NSOS ``JambaModel`` on top of ``nsos_ext`` and
+standardizes device selection, streaming incremental decoding and small
+supervised datasets.
+"""
+
+from .core import (
+    DecodeBenchmark,
+    GenerationResult,
+    MambaModuleConfig,
+    NSOSMamba,
+    detect_device,
+)
+from .datasets import (
+    CharTokenizer,
+    TextPair,
+    load_text_pairs,
+    load_text_pairs_json,
+    load_text_pairs_jsonl,
+    pairs_to_token_ids,
+)
+
+__all__ = [
+    "CharTokenizer",
+    "DecodeBenchmark",
+    "GenerationResult",
+    "MambaModuleConfig",
+    "NSOSMamba",
+    "TextPair",
+    "detect_device",
+    "load_text_pairs",
+    "load_text_pairs_json",
+    "load_text_pairs_jsonl",
+    "pairs_to_token_ids",
+]

nsos_mamba-0.1.0/nsos_mamba/core.py

@@ -0,0 +1,465 @@
+from __future__ import annotations
+
+import os
+import random
+import time
+from dataclasses import dataclass, field
+from typing import Callable, Iterable, List, Optional, Sequence, Tuple
+
+import numpy as np
+
+from .datasets import CharTokenizer, TextPair, pairs_to_token_ids
+
+
+def _load_nsos_ext():
+    ext_path = os.environ.get("NSOS_EXT_PATH")
+    if ext_path:
+        import sys
+
+        sys.path.insert(0, ext_path)
+    try:
+        import nsos_ext as nsos  # type: ignore
+    except ImportError as exc:
+        raise ImportError(
+            "nsos_ext is not importable. Add the built nsos_ext directory to "
+            "PYTHONPATH/sys.path before importing nsos_mamba."
+        ) from exc
+    return nsos
+
+
+def _torch_cuda_available() -> bool:
+    try:
+        import torch  # type: ignore
+
+        return bool(torch.cuda.is_available())
+    except Exception:
+        return False
+
+
+def detect_device(nsos_ext=None, requested: Optional[str] = None):
+    """Resolve CPU/GPU from a flag or ``NSOS_MAMBA_DEVICE``.
+
+    Accepted values: ``auto`` (default), ``gpu``/``cuda`` and ``cpu``.
+    ``auto`` selects GPU only when PyTorch can see CUDA; construction still has
+    a CPU fallback in ``NSOSMamba`` in case the native extension cannot allocate
+    on the selected device.
+    """
+
+    nsos = nsos_ext or _load_nsos_ext()
+    mode = (requested or os.environ.get("NSOS_MAMBA_DEVICE", "auto")).strip().lower()
+    if mode in ("cpu", "host"):
+        return nsos.Device.CPU
+    if mode in ("gpu", "cuda"):
+        return nsos.Device.GPU
+    if mode not in ("", "auto"):
+        raise ValueError("device must be one of: auto, gpu, cuda, cpu")
+    return nsos.Device.GPU if _torch_cuda_available() else nsos.Device.CPU
+
+
+def _env_bool(name: str, default: bool) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return default
+    return value.strip().lower() not in ("0", "false", "no", "off")
+
+
+def _safe_set(obj, name: str, value) -> None:
+    try:
+        setattr(obj, name, value)
+    except Exception:
+        pass
+
+
+@dataclass
+class MambaModuleConfig:
+    """Configuration for the standalone faithful-Mamba wrapper."""
+
+    vocab_size: int
+    num_layers: int = 12
+    d_model: int = 128
+    d_state: int = 64
+    max_context_tokens: int = 4096
+    n_heads: int = 4
+    n_kv_heads: int = 2
+    mamba_expand: int = 2
+    mamba_head_dim: int = 64
+    mamba_n_groups: int = 1
+    tie_word_embeddings: bool = False
+    seed: Optional[int] = 42
+    device: str = "auto"
+    streaming: bool = True
+    allow_cpu_fallback: bool = True
+
+    @classmethod
+    def from_env(cls, vocab_size: int) -> "MambaModuleConfig":
+        return cls(
+            vocab_size=vocab_size,
+            num_layers=int(os.environ.get("NSOS_MAMBA_LAYERS", "12")),
+            d_model=int(os.environ.get("NSOS_MAMBA_DMODEL", "128")),
+            d_state=int(os.environ.get("NSOS_MAMBA_DSTATE", "64")),
+            max_context_tokens=int(os.environ.get("NSOS_MAMBA_CONTEXT", "4096")),
+            device=os.environ.get("NSOS_MAMBA_DEVICE", "auto"),
+            streaming=_env_bool("NSOS_MAMBA_STREAMING", True),
+        )
+
+
+@dataclass
+class GenerationResult:
+    prompt_ids: List[int]
+    generated_ids: List[int]
+    text: Optional[str]
+    prefill_ms: float
+    decode_ms: float
+    ttft_ms: float
+    first_incremental_ms: float
+    generated_tokens: int
+    decode_tokens_per_sec: float
+    end_to_end_tokens_per_sec: float
+    streaming: bool
+
+
+@dataclass
+class DecodeBenchmark:
+    runs: List[GenerationResult] = field(default_factory=list)
+
+    @property
+    def avg_decode_tokens_per_sec(self) -> float:
+        return float(np.mean([r.decode_tokens_per_sec for r in self.runs])) if self.runs else 0.0
+
+    @property
+    def avg_end_to_end_tokens_per_sec(self) -> float:
+        return float(np.mean([r.end_to_end_tokens_per_sec for r in self.runs])) if self.runs else 0.0
+
+    @property
+    def avg_prefill_ms(self) -> float:
+        return float(np.mean([r.prefill_ms for r in self.runs])) if self.runs else 0.0
+
+    @property
+    def avg_first_incremental_ms(self) -> float:
+        return float(np.mean([r.first_incremental_ms for r in self.runs])) if self.runs else 0.0
+
+    def print(self) -> None:
+        print("\n========== NSOS MAMBA DECODE BENCH ==========")
+        for idx, run in enumerate(self.runs):
+            print(
+                f"run {idx}: {run.generated_tokens} tok | "
+                f"decode={run.decode_ms / 1000.0:.4f}s | "
+                f"decode_tok/s={run.decode_tokens_per_sec:.2f} | "
+                f"e2e_tok/s={run.end_to_end_tokens_per_sec:.2f} | "
+                f"prefill={run.prefill_ms:.2f} ms | "
+                f"first_incremental={run.first_incremental_ms:.2f} ms"
+            )
+        print("---------------------------------------------")
+        print("avg decode tok/s:", self.avg_decode_tokens_per_sec)
+        print("avg e2e tok/s:", self.avg_end_to_end_tokens_per_sec)
+        print("avg prefill ms:", self.avg_prefill_ms)
+        print("avg first incremental ms:", self.avg_first_incremental_ms)
+
+
+class NSOSMamba:
+    """Pure faithful-Mamba NSOS module with streaming incremental by default."""
+
+    def __init__(self, config: MambaModuleConfig, nsos_ext=None):
+        self.nsos = nsos_ext or _load_nsos_ext()
+        self.config = config
+        if config.seed is not None:
+            self.nsos.set_seed(int(config.seed))
+
+        requested_device = detect_device(self.nsos, config.device)
+        self.device = requested_device
+        try:
+            self.model = self.nsos.JambaModel(self._make_native_config(), requested_device)
+            self.model.to(requested_device)
+        except Exception:
+            if not config.allow_cpu_fallback or requested_device == self.nsos.Device.CPU:
+                raise
+            self.device = self.nsos.Device.CPU
+            self.model = self.nsos.JambaModel(self._make_native_config(), self.device)
+            self.model.to(self.device)
+
+        self.set_streaming(config.streaming)
+
+    def _make_native_config(self):
+        cfg = self.nsos.ModelConfig()
+        cfg.num_layers = int(self.config.num_layers)
+        cfg.d_model = int(self.config.d_model)
+        cfg.vocab_size = int(self.config.vocab_size)
+        cfg.n_heads = int(self.config.n_heads)
+        cfg.n_kv_heads = int(self.config.n_kv_heads)
+        cfg.max_context_tokens = int(self.config.max_context_tokens)
+
+        cfg.use_moe = False
+        cfg.use_kan = False
+        cfg.use_ttt = False
+        cfg.dropout = 0.0
+
+        cfg.mamba2_faithful = True
+        cfg.mamba_expand = int(self.config.mamba_expand)
+        cfg.mamba_head_dim = int(self.config.mamba_head_dim)
+        cfg.tie_word_embeddings = bool(self.config.tie_word_embeddings)
+
+        _safe_set(cfg, "mamba_state_expansion", int(self.config.d_state))
+        _safe_set(cfg, "mamba_n_groups", int(self.config.mamba_n_groups))
+
+        # Disable attention by putting the slot outside any real layer schedule.
+        cfg.attention_period = 1_000_000
+        cfg.attention_slot = 0
+        return cfg
+
+    @property
+    def vocab_size(self) -> int:
+        return int(self.config.vocab_size)
+
+    def set_streaming(self, enabled: bool = True) -> bool:
+        try:
+            if enabled and not self.model.supports_streaming_inference():
+                self.model.set_streaming_inference(False)
+                return False
+            self.model.set_streaming_inference(bool(enabled))
+            return bool(enabled)
+        except Exception:
+            return False
+
+    def _logits_numpy(self, tensor) -> np.ndarray:
+        return np.asarray(tensor.cpu().numpy()).reshape(-1, self.vocab_size)
+
+    def fit_token_pairs(
+        self,
+        pairs: Sequence[Tuple[Sequence[int], Sequence[int]]],
+        *,
+        steps: int = 1200,
+        batch_size: int = 16,
+        learning_rate: float = 3e-3,
+        qat: bool = False,
+        seed: int = 123,
+        progress_every: Optional[int] = None,
+        callback: Optional[Callable[[int, float], None]] = None,
+    ) -> List[float]:
+        if not pairs:
+            raise ValueError("fit_token_pairs requires at least one pair")
+
+        trainer = self.nsos.Trainer(self.model, float(learning_rate))
+        trainer.warmup_steps = max(1, min(50, steps // 10 if steps > 10 else 1))
+        trainer.total_training_steps = int(steps)
+        trainer.first_token_loss_scale = 1.0
+        trainer.eos_loss_scale = 1.0
+        trainer.phase_scheduler.progressive_qat_enabled = bool(qat)
+
+        rng = random.Random(seed)
+        self.model.set_training_mode(True)
+        self.set_streaming(False)
+
+        losses: List[float] = []
+        window = progress_every or max(1, steps // 6)
+        running = 0.0
+        for step in range(int(steps)):
+            prompts: List[List[int]] = []
+            answers: List[List[int]] = []
+            for _ in range(int(batch_size)):
+                prompt, answer = rng.choice(pairs)
+                prompts.append(list(prompt))
+                answers.append(list(answer))
+            loss = float(trainer.train_supervised_batch(prompts, answers))
+            losses.append(loss)
+            running += loss
+            if callback and (step + 1) % window == 0:
+                callback(step + 1, running / float(window))
+                running = 0.0
+
+        self.model.set_training_mode(False)
+        self.set_streaming(self.config.streaming)
+        return losses
+
+    def fit_text_pairs(
+        self,
+        pairs: Sequence[TextPair],
+        tokenizer: CharTokenizer,
+        **kwargs,
+    ) -> List[float]:
+        token_pairs = pairs_to_token_ids(pairs, tokenizer)
+        return self.fit_token_pairs(token_pairs, **kwargs)
+
+    def generate_ids(
+        self,
+        prompt_ids: Sequence[int],
+        *,
+        max_new_tokens: int = 128,
+        eos_token_id: Optional[int] = None,
+        stop_on_eos: bool = True,
+        streaming: Optional[bool] = None,
+        tokenizer: Optional[CharTokenizer] = None,
+    ) -> GenerationResult:
+        if max_new_tokens <= 0:
+            return GenerationResult(
+                prompt_ids=list(prompt_ids),
+                generated_ids=[],
+                text="",
+                prefill_ms=0.0,
+                decode_ms=0.0,
+                ttft_ms=0.0,
+                first_incremental_ms=0.0,
+                generated_tokens=0,
+                decode_tokens_per_sec=0.0,
+                end_to_end_tokens_per_sec=0.0,
+                streaming=bool(streaming),
+            )
+
+        use_streaming = self.config.streaming if streaming is None else bool(streaming)
+        self.model.set_training_mode(False)
+        self.set_streaming(use_streaming)
+
+        if use_streaming:
+            return self._generate_ids_streaming(
+                prompt_ids,
+                max_new_tokens=max_new_tokens,
+                eos_token_id=eos_token_id,
+                stop_on_eos=stop_on_eos,
+                tokenizer=tokenizer,
+            )
+        return self._generate_ids_eager(
+            prompt_ids,
+            max_new_tokens=max_new_tokens,
+            eos_token_id=eos_token_id,
+            stop_on_eos=stop_on_eos,
+            tokenizer=tokenizer,
+        )
+
+    def _generate_ids_streaming(
+        self,
+        prompt_ids: Sequence[int],
+        *,
+        max_new_tokens: int,
+        eos_token_id: Optional[int],
+        stop_on_eos: bool,
+        tokenizer: Optional[CharTokenizer],
+    ) -> GenerationResult:
+        self.model.reset_session()
+        prompt = list(prompt_ids)
+
+        prefill_t0 = time.perf_counter()
+        logits = self._logits_numpy(self.model.forward_ids(prompt))
+        prefill_s = time.perf_counter() - prefill_t0
+
+        next_id = int(np.argmax(logits[-1]))
+        generated = [next_id]
+        if stop_on_eos and eos_token_id is not None and next_id == eos_token_id:
+            text = tokenizer.decode(generated) if tokenizer else None
+            return self._result(prompt, generated, text, prefill_s, 0.0, 0.0, True)
+
+        decode_t0 = time.perf_counter()
+        first_incremental_s = 0.0
+        for idx in range(1, max_new_tokens):
+            step_t0 = time.perf_counter()
+            logits = self._logits_numpy(self.model.forward_ids([next_id]))
+            if idx == 1:
+                first_incremental_s = time.perf_counter() - step_t0
+            next_id = int(np.argmax(logits[-1]))
+            generated.append(next_id)
+            if stop_on_eos and eos_token_id is not None and next_id == eos_token_id:
+                break
+        decode_s = time.perf_counter() - decode_t0
+        text = tokenizer.decode(generated) if tokenizer else None
+        return self._result(prompt, generated, text, prefill_s, decode_s, first_incremental_s, True)
+
+    def _generate_ids_eager(
+        self,
+        prompt_ids: Sequence[int],
+        *,
+        max_new_tokens: int,
+        eos_token_id: Optional[int],
+        stop_on_eos: bool,
+        tokenizer: Optional[CharTokenizer],
+    ) -> GenerationResult:
+        prompt = list(prompt_ids)
+        seq = list(prompt)
+        generated: List[int] = []
+        t0 = time.perf_counter()
+        first_s = 0.0
+        for idx in range(max_new_tokens):
+            step_t0 = time.perf_counter()
+            logits = self._logits_numpy(self.model.forward_ids(seq))
+            if idx == 0:
+                first_s = time.perf_counter() - step_t0
+            next_id = int(np.argmax(logits[-1]))
+            generated.append(next_id)
+            seq.append(next_id)
+            if stop_on_eos and eos_token_id is not None and next_id == eos_token_id:
+                break
+        total_s = time.perf_counter() - t0
+        text = tokenizer.decode(generated) if tokenizer else None
+        return self._result(prompt, generated, text, first_s, max(total_s - first_s, 0.0), first_s, False)
+
+    def _result(
+        self,
+        prompt: List[int],
+        generated: List[int],
+        text: Optional[str],
+        prefill_s: float,
+        decode_s: float,
+        first_incremental_s: float,
+        streaming: bool,
+    ) -> GenerationResult:
+        # In streaming mode the first generated token comes from prefill; steady
+        # decode throughput counts only tokens produced after prefill.
+        steady_tokens = max(len(generated) - 1, 0) if streaming else len(generated)
+        decode_tps = steady_tokens / max(decode_s, 1e-9)
+        e2e_tps = len(generated) / max(prefill_s + decode_s, 1e-9)
+        return GenerationResult(
+            prompt_ids=prompt,
+            generated_ids=generated,
+            text=text,
+            prefill_ms=prefill_s * 1000.0,
+            decode_ms=decode_s * 1000.0,
+            ttft_ms=prefill_s * 1000.0,
+            first_incremental_ms=first_incremental_s * 1000.0,
+            generated_tokens=len(generated),
+            decode_tokens_per_sec=decode_tps,
+            end_to_end_tokens_per_sec=e2e_tps,
+            streaming=streaming,
+        )
+
+    def generate_text(
+        self,
+        prompt: str,
+        tokenizer: CharTokenizer,
+        *,
+        max_new_tokens: int = 128,
+        stop_on_eos: bool = True,
+    ) -> GenerationResult:
+        prompt_ids = tokenizer.encode(prompt, bos=True, eos=False)
+        return self.generate_ids(
+            prompt_ids,
+            max_new_tokens=max_new_tokens,
+            eos_token_id=tokenizer.eos_id,
+            stop_on_eos=stop_on_eos,
+            tokenizer=tokenizer,
+        )
+
+    def benchmark_decode(
+        self,
+        prompt_ids: Sequence[int],
+        *,
+        max_new_tokens: int = 256,
+        repeats: int = 5,
+        warmup_tokens: int = 8,
+        eos_token_id: Optional[int] = None,
+    ) -> DecodeBenchmark:
+        # Warmup: one short streaming run, no EOS stop.
+        self.generate_ids(
+            prompt_ids,
+            max_new_tokens=max(1, warmup_tokens),
+            eos_token_id=eos_token_id,
+            stop_on_eos=False,
+            streaming=True,
+        )
+        runs = [
+            self.generate_ids(
+                prompt_ids,
+                max_new_tokens=max_new_tokens,
+                eos_token_id=eos_token_id,
+                stop_on_eos=False,
+                streaming=True,
+            )
+            for _ in range(repeats)
+        ]
+        return DecodeBenchmark(runs)

nsos_mamba-0.1.0/nsos_mamba/datasets.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, List, Mapping, Sequence, Tuple
+
+
+@dataclass(frozen=True)
+class TextPair:
+    """One supervised prompt -> answer example."""
+
+    prompt: str
+    answer: str
+
+
+class CharTokenizer:
+    """Small deterministic char-level tokenizer for demos and probes.
+
+    It is deliberately simple.  Production training should use the project's
+    tokenizer pack, but this class is enough for synthetic tasks such as:
+    "Quem é você?" -> "Oxta".
+    """
+
+    pad_token = "<PAD>"
+    bos_token = "<BOS>"
+    eos_token = "<EOS>"
+
+    def __init__(self, tokens: Sequence[str]):
+        unique = list(dict.fromkeys(tokens))
+        required = [self.pad_token, self.bos_token, self.eos_token]
+        merged = required + [t for t in unique if t not in required]
+        self.tokens: List[str] = merged
+        self.stoi = {tok: idx for idx, tok in enumerate(self.tokens)}
+        self.itos = {idx: tok for tok, idx in self.stoi.items()}
+
+    @classmethod
+    def from_texts(cls, texts: Iterable[str]) -> "CharTokenizer":
+        chars = sorted({ch for text in texts for ch in text})
+        return cls([cls.pad_token, cls.bos_token, cls.eos_token, *chars])
+
+    @property
+    def pad_id(self) -> int:
+        return self.stoi[self.pad_token]
+
+    @property
+    def bos_id(self) -> int:
+        return self.stoi[self.bos_token]
+
+    @property
+    def eos_id(self) -> int:
+        return self.stoi[self.eos_token]
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.tokens)
+
+    def encode(self, text: str, *, bos: bool = False, eos: bool = False) -> List[int]:
+        ids: List[int] = []
+        if bos:
+            ids.append(self.bos_id)
+        for ch in text:
+            if ch not in self.stoi:
+                raise KeyError(f"character not in tokenizer vocabulary: {ch!r}")
+            ids.append(self.stoi[ch])
+        if eos:
+            ids.append(self.eos_id)
+        return ids
+
+    def decode(self, ids: Iterable[int], *, stop_at_eos: bool = True) -> str:
+        out: List[str] = []
+        for token_id in ids:
+            token_id = int(token_id)
+            if token_id == self.eos_id and stop_at_eos:
+                break
+            if token_id in (self.pad_id, self.bos_id, self.eos_id):
+                continue
+            out.append(self.itos.get(token_id, ""))
+        return "".join(out)
+
+
+def load_text_pairs(pairs: Iterable[Tuple[str, str] | Mapping[str, str] | TextPair]) -> List[TextPair]:
+    """Normalize in-memory examples into ``TextPair`` objects."""
+
+    normalized: List[TextPair] = []
+    for item in pairs:
+        if isinstance(item, TextPair):
+            normalized.append(item)
+        elif isinstance(item, Mapping):
+            normalized.append(TextPair(str(item["prompt"]), str(item["answer"])))
+        else:
+            prompt, answer = item
+            normalized.append(TextPair(str(prompt), str(answer)))
+    return normalized
+
+
+def load_text_pairs_json(
+    path: str | Path,
+    *,
+    prompt_key: str = "prompt",
+    answer_key: str = "answer",
+) -> List[TextPair]:
+    """Load examples from a JSON file.
+
+    Accepted shapes:
+
+    - ``[{"prompt": "...", "answer": "..."}]``
+    - ``{"data": [{"prompt": "...", "answer": "..."}]}``
+    - ``{"examples": [{"prompt": "...", "answer": "..."}]}``
+    """
+
+    raw = json.loads(Path(path).read_text(encoding="utf-8"))
+    if isinstance(raw, Mapping):
+        if "data" in raw:
+            raw = raw["data"]
+        elif "examples" in raw:
+            raw = raw["examples"]
+    if not isinstance(raw, list):
+        raise ValueError("JSON dataset must be a list or an object with data/examples")
+    return [
+        TextPair(prompt=str(row[prompt_key]), answer=str(row[answer_key]))
+        for row in raw
+    ]
+
+
+def load_text_pairs_jsonl(
+    path: str | Path,
+    *,
+    prompt_key: str = "prompt",
+    answer_key: str = "answer",
+) -> List[TextPair]:
+    """Load one ``{"prompt": "...", "answer": "..."}`` object per line."""
+
+    examples: List[TextPair] = []
+    for line_no, line in enumerate(Path(path).read_text(encoding="utf-8").splitlines(), 1):
+        stripped = line.strip()
+        if not stripped:
+            continue
+        row = json.loads(stripped)
+        if not isinstance(row, Mapping):
+            raise ValueError(f"JSONL line {line_no} is not an object")
+        examples.append(TextPair(prompt=str(row[prompt_key]), answer=str(row[answer_key])))
+    return examples
+
+
+def pairs_to_token_ids(
+    pairs: Sequence[TextPair],
+    tokenizer: CharTokenizer,
+    *,
+    add_bos: bool = True,
+    add_eos: bool = True,
+) -> List[Tuple[List[int], List[int]]]:
+    """Convert text pairs to NSOS supervised prompt/answer token lists."""
+
+    return [
+        (
+            tokenizer.encode(pair.prompt, bos=add_bos, eos=False),
+            tokenizer.encode(pair.answer, bos=False, eos=add_eos),
+        )
+        for pair in pairs
+    ]

nsos_mamba-0.1.0/nsos_mamba.egg-info/PKG-INFO

@@ -0,0 +1,228 @@
+Metadata-Version: 2.4
+Name: nsos-mamba
+Version: 0.1.0
+Summary: Wrapper Python isolado para o módulo Mamba do NSOS com streaming incremental padrão.
+Author: Oxta
+Project-URL: Homepage, https://github.com/vitorGgC569/reimagined
+Project-URL: Documentation, https://github.com/vitorGgC569/reimagined/tree/nsos-gpu-phases12/OXN/nsos/docs
+Project-URL: Source, https://github.com/vitorGgC569/reimagined
+Keywords: mamba,ssm,nsos,edge-ai,streaming
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.23
+
+# NSOS Mamba
+
+`nsos-mamba` é um wrapper Python simples para usar apenas o módulo Mamba do
+NSOS. Ele configura automaticamente o modelo em modo Mamba-2 faithful, ativa
+streaming incremental por padrão e oferece utilitários pequenos para treinar
+pares `prompt -> resposta`.
+
+Este pacote não instala o runtime nativo `nsos_ext`. Antes de usar, compile ou
+copie o `nsos_ext.so`/`nsos_ext.pyd` e deixe-o acessível via `PYTHONPATH` ou
+pela variável `NSOS_EXT_PATH`.
+
+## Instalação
+
+```bash
+pip install nsos-mamba
+```
+
+Se o `nsos_ext` estiver em uma pasta separada:
+
+```bash
+export NSOS_EXT_PATH=/caminho/para/nsos_ext
+```
+
+No Colab:
+
+```python
+import sys, os
+
+sys.path.insert(0, "/content/nsos_ext_mambavs")
+# ou:
+os.environ["NSOS_EXT_PATH"] = "/content/nsos_ext_mambavs"
+```
+
+## Configurações principais
+
+```python
+from nsos_mamba import MambaModuleConfig
+
+cfg = MambaModuleConfig(
+    vocab_size=128,
+    num_layers=12,
+    d_model=128,
+    device="auto",
+    streaming=True,
+)
+```
+
+Campos mais usados:
+
+- `vocab_size`: tamanho do vocabulário.
+- `num_layers`: número de camadas Mamba.
+- `d_model`: largura do modelo.
+- `d_state`: tamanho do estado SSM. Padrão: `64`.
+- `max_context_tokens`: limite de contexto. Padrão: `4096`.
+- `device`: `"auto"`, `"gpu"`, `"cuda"` ou `"cpu"`.
+- `streaming`: ativa decode incremental. Padrão: `True`.
+- `allow_cpu_fallback`: se a GPU falhar, cai para CPU. Padrão: `True`.
+
+Variáveis de ambiente equivalentes:
+
+```bash
+NSOS_MAMBA_DEVICE=auto
+NSOS_MAMBA_STREAMING=1
+NSOS_MAMBA_LAYERS=12
+NSOS_MAMBA_DMODEL=128
+NSOS_MAMBA_DSTATE=64
+NSOS_MAMBA_CONTEXT=4096
+```
+
+## Exemplo mínimo: "Quem é você?" -> "Oxta"
+
+```python
+from nsos_mamba import CharTokenizer, MambaModuleConfig, NSOSMamba, TextPair
+
+pairs = [
+    TextPair("Quem é você?", "Oxta"),
+    TextPair("quem é você?", "Oxta"),
+    TextPair("Quem e voce?", "Oxta"),
+]
+
+tok = CharTokenizer.from_texts(
+    [p.prompt for p in pairs] + [p.answer for p in pairs]
+)
+
+cfg = MambaModuleConfig(
+    vocab_size=tok.vocab_size,
+    num_layers=12,
+    d_model=128,
+    device="auto",
+    streaming=True,
+)
+
+mamba = NSOSMamba(cfg)
+
+mamba.fit_text_pairs(
+    pairs,
+    tok,
+    steps=1200,
+    batch_size=16,
+    learning_rate=3e-3,
+    callback=lambda step, loss: print(f"[train] step {step} loss~{loss:.4f}"),
+)
+
+result = mamba.generate_text("Quem é você?", tok, max_new_tokens=8)
+
+print("resposta:", repr(result.text))
+print("tokens/s:", result.decode_tokens_per_sec)
+```
+
+Saída esperada:
+
+```text
+resposta: 'Oxta'
+```
+
+## Streaming incremental
+
+O streaming incremental é o padrão. O fluxo usado é:
+
+1. processa o prompt uma vez (`prefill`);
+2. gera cada token novo passando apenas o último token;
+3. reaproveita o estado interno do modelo.
+
+Benchmark:
+
+```python
+prompt_ids = tok.encode("Quem é você?", bos=True)
+bench = mamba.benchmark_decode(prompt_ids, max_new_tokens=1024, repeats=5)
+bench.print()
+```
+
+Métricas:
+
+- `prefill_ms`: tempo para processar o prompt.
+- `first_incremental_ms`: primeiro passo incremental após o prefill.
+- `decode_tokens_per_sec`: tokens/s do decode incremental.
+- `end_to_end_tokens_per_sec`: tokens/s incluindo prefill.
+
+## Datasets
+
+### Lista em memória
+
+```python
+from nsos_mamba import TextPair
+
+pairs = [
+    TextPair(prompt="Quem é você?", answer="Oxta"),
+    TextPair(prompt="Qual seu nome?", answer="Oxta"),
+]
+```
+
+### JSON
+
+Arquivo `dataset.json`:
+
+```json
+[
+  {"prompt": "Quem é você?", "answer": "Oxta"},
+  {"prompt": "Qual seu nome?", "answer": "Oxta"}
+]
+```
+
+Uso:
+
+```python
+from nsos_mamba import load_text_pairs_json
+
+pairs = load_text_pairs_json("dataset.json")
+```
+
+Também aceita:
+
+```json
+{"data": [{"prompt": "...", "answer": "..."}]}
+```
+
+ou:
+
+```json
+{"examples": [{"prompt": "...", "answer": "..."}]}
+```
+
+### JSONL
+
+Arquivo `dataset.jsonl`:
+
+```jsonl
+{"prompt": "Quem é você?", "answer": "Oxta"}
+{"prompt": "Qual seu nome?", "answer": "Oxta"}
+```
+
+Uso:
+
+```python
+from nsos_mamba import load_text_pairs_jsonl
+
+pairs = load_text_pairs_jsonl("dataset.jsonl")
+```
+
+## Observações
+
+- `CharTokenizer` é para testes, demos e dados sintéticos pequenos.
+- Para treino real, use um tokenizer próprio do projeto.
+- O pacote expõe o Mamba isolado; não ativa Attention, MoE, KAN, TTT, CHRASS,
+  MCTS ou memória externa.
+- O pacote é uma camada de uso sobre `nsos_ext`, não uma reimplementação em
+  Python.

nsos_mamba-0.1.0/nsos_mamba.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+README.md
+pyproject.toml
+nsos_mamba/__init__.py
+nsos_mamba/core.py
+nsos_mamba/datasets.py
+nsos_mamba.egg-info/PKG-INFO
+nsos_mamba.egg-info/SOURCES.txt
+nsos_mamba.egg-info/dependency_links.txt
+nsos_mamba.egg-info/requires.txt
+nsos_mamba.egg-info/top_level.txt

nsos_mamba-0.1.0/nsos_mamba.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

nsos_mamba-0.1.0/nsos_mamba.egg-info/requires.txt

@@ -0,0 +1 @@
+numpy>=1.23

nsos_mamba-0.1.0/nsos_mamba.egg-info/top_level.txt

@@ -0,0 +1 @@
+nsos_mamba

nsos_mamba-0.1.0/pyproject.toml

@@ -0,0 +1,35 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "nsos-mamba"
+version = "0.1.0"
+description = "Wrapper Python isolado para o módulo Mamba do NSOS com streaming incremental padrão."
+readme = "README.md"
+requires-python = ">=3.10"
+authors = [
+  { name = "Oxta" }
+]
+keywords = ["mamba", "ssm", "nsos", "edge-ai", "streaming"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+  "numpy>=1.23",
+]
+
+[project.urls]
+Homepage = "https://github.com/vitorGgC569/reimagined"
+Documentation = "https://github.com/vitorGgC569/reimagined/tree/nsos-gpu-phases12/OXN/nsos/docs"
+Source = "https://github.com/vitorGgC569/reimagined"
+
+[tool.setuptools]
+packages = ["nsos_mamba"]

nsos_mamba-0.1.0/setup.cfg

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