bit-ttt-engine 0.6.2__cp310-cp310-win_amd64.whl → 0.7.0__cp310-cp310-win_amd64.whl

cortex_rust/cli.py

@@ -0,0 +1,381 @@
+"""Bit-TTT-Engine CLI — interactive chat and text generation.
+
+Usage:
+    python -m cortex_rust chat model.gguf [--template llama3] [--system "..."]
+    python -m cortex_rust generate model.gguf --prompt "..." [--max-tokens 200]
+    python -m cortex_rust info model.gguf
+    
+Or if installed with pip:
+    bit-ttt chat model.gguf
+    bit-ttt generate model.gguf --prompt "Hello"
+    bit-ttt info model.gguf
+"""
+
+import argparse
+import sys
+import os
+import time
+from typing import Optional
+
+
+def _resolve_model_path(model_path: str) -> str:
+    """Resolve model path: local file or HuggingFace repo ID.
+    
+    If model_path is a local file, return as-is.
+    If it looks like a HF repo ID (contains '/'), auto-download the best GGUF.
+    
+    Examples:
+        "model.gguf"                              → local file
+        "bartowski/Llama-3-8B-Instruct-GGUF"      → download Q4_K_M
+        "user/repo:Q8_0"                          → download Q8_0 variant
+    """
+    # Local file exists — use directly
+    if os.path.exists(model_path):
+        return model_path
+
+    # Not a repo ID pattern — treat as local path (will error later)
+    if "/" not in model_path:
+        return model_path
+
+    # Parse repo ID and optional quantization hint
+    # Format: "user/repo" or "user/repo:Q8_0"
+    if ":" in model_path:
+        repo_id, quant_hint = model_path.rsplit(":", 1)
+    else:
+        repo_id, quant_hint = model_path, "Q4_K_M"
+
+    print(f"📥 Resolving model from HuggingFace: {repo_id}")
+
+    try:
+        from huggingface_hub import HfApi, hf_hub_download
+    except ImportError:
+        print("❌ huggingface_hub required for auto-download.")
+        print("   Install: pip install huggingface_hub")
+        sys.exit(1)
+
+    cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "bit-ttt")
+    os.makedirs(cache_dir, exist_ok=True)
+
+    api = HfApi()
+
+    try:
+        # List GGUF files in the repo
+        files = api.list_repo_files(repo_id)
+        gguf_files = [f for f in files if f.endswith(".gguf")]
+
+        if not gguf_files:
+            print(f"❌ No GGUF files found in {repo_id}")
+            sys.exit(1)
+
+        # Find best match for quantization hint
+        quant_lower = quant_hint.lower().replace("-", "_")
+        matched = [f for f in gguf_files if quant_lower in f.lower()]
+
+        if matched:
+            target = matched[0]
+        else:
+            # Prefer Q4_K_M > Q4_K_S > first available
+            for pref in ["Q4_K_M", "Q4_K_S", "Q4_0", "Q5_K_M", "Q8_0"]:
+                pref_match = [f for f in gguf_files if pref.lower() in f.lower()]
+                if pref_match:
+                    target = pref_match[0]
+                    break
+            else:
+                target = gguf_files[0]
+
+        print(f"📦 Selected: {target}")
+        print(f"   Downloading to cache...")
+
+        local_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=target,
+            cache_dir=cache_dir,
+            local_dir=os.path.join(cache_dir, repo_id.replace("/", "--")),
+        )
+
+        # Also try to download tokenizer.json
+        if "tokenizer.json" in files:
+            try:
+                hf_hub_download(
+                    repo_id=repo_id,
+                    filename="tokenizer.json",
+                    cache_dir=cache_dir,
+                    local_dir=os.path.join(cache_dir, repo_id.replace("/", "--")),
+                )
+            except Exception:
+                pass  # tokenizer is optional
+
+        print(f"✅ Downloaded: {local_path}")
+        return local_path
+
+    except Exception as e:
+        print(f"❌ Failed to download from {repo_id}: {e}")
+        sys.exit(1)
+
+
+def _load_model(model_path: str, device: str = "cuda", tokenizer: Optional[str] = None):
+    """Load a GGUF model, auto-detecting tokenizer if needed.
+    
+    model_path can be a local file or a HuggingFace repo ID.
+    """
+    from cortex_rust import QGgufModel
+
+    # Resolve HF repo ID → local path
+    resolved_path = _resolve_model_path(model_path)
+    
+    tok = tokenizer
+    if tok is None:
+        # Try to find tokenizer.json in same directory
+        model_dir = os.path.dirname(os.path.abspath(resolved_path))
+        candidate = os.path.join(model_dir, "tokenizer.json")
+        if os.path.exists(candidate):
+            tok = candidate
+    
+    kwargs = {"device": device}
+    if tok:
+        kwargs["tokenizer"] = tok
+    
+    return QGgufModel(resolved_path, **kwargs)
+
+
+def _apply_feature_flags(model, args):
+    """Apply --lora, --q8-cache, --ttt flags to a loaded model."""
+    if getattr(args, 'q8_cache', False):
+        model.enable_q8_kv_cache(True)
+        print("   🔧 Q8 KV cache enabled")
+    if getattr(args, 'ttt', False):
+        model.enable_ttt(True)
+        print("   🔧 TTT (Test-Time Training) enabled")
+    if getattr(args, 'lora', None):
+        model.load_lora(args.lora)
+        model.enable_lora(True)
+        print(f"   🔧 LoRA loaded: {args.lora}")
+    return model
+
+
+def cmd_info(args):
+    """Show model information."""
+    model = _load_model(args.model, device="cpu")
+    config = model.config
+    
+    print(f"📦 Model: {os.path.basename(args.model)}")
+    print(f"   Architecture:  {getattr(config, 'arch', 'unknown')}")
+    print(f"   Vocab size:    {config.vocab_size:,}")
+    print(f"   Hidden dim:    {config.hidden_dim:,}")
+    print(f"   Layers:        {config.num_layers}")
+    print(f"   Q heads:       {config.n_heads}")
+    print(f"   KV heads:      {config.n_kv_heads}")
+    print(f"   Head dim:      {config.hidden_dim // config.n_heads}")
+    print(f"   GQA ratio:     {config.n_heads // config.n_kv_heads}")
+    
+    from cortex_rust.chat import detect_template
+    template = detect_template(args.model)
+    print(f"   Chat template: {template}")
+
+
+def cmd_generate(args):
+    """Generate text from a prompt."""
+    from cortex_rust.chat import format_simple, detect_template
+    
+    print(f"Loading {os.path.basename(args.model)}...", flush=True)
+    model = _load_model(args.model, device=args.device, tokenizer=args.tokenizer)
+    _apply_feature_flags(model, args)
+    
+    template = args.template or detect_template(args.model)
+    
+    if args.raw:
+        prompt = args.prompt
+    else:
+        prompt = format_simple(args.prompt, system_message=args.system, template=template)
+    
+    print(f"Generating (template={template}, max_tokens={args.max_tokens})...\n")
+    
+    start = time.time()
+    
+    # Use streaming callback
+    token_count = [0]
+    def on_token(token_str):
+        print(token_str, end="", flush=True)
+        token_count[0] += 1
+    
+    model.generate_with_callback(
+        prompt,
+        on_token,
+        max_tokens=args.max_tokens,
+        temperature=args.temperature,
+        top_k=args.top_k,
+        top_p=args.top_p,
+        repetition_penalty=args.repetition_penalty,
+    )
+    
+    elapsed = time.time() - start
+    speed = token_count[0] / elapsed if elapsed > 0 else 0
+    print(f"\n\n--- {token_count[0]} tokens in {elapsed:.1f}s ({speed:.1f} tok/s) ---")
+
+
+def cmd_chat(args):
+    """Interactive chat mode."""
+    from cortex_rust.chat import format_chat, detect_template
+    
+    print(f"Loading {os.path.basename(args.model)}...", flush=True)
+    model = _load_model(args.model, device=args.device, tokenizer=args.tokenizer)
+    _apply_feature_flags(model, args)
+    
+    template = args.template or detect_template(args.model)
+    
+    print(f"💬 Chat mode (template={template})")
+    print(f"   Type 'quit' or Ctrl+C to exit")
+    print(f"   Type '/reset' to clear conversation")
+    print(f"   Type '/system <msg>' to change system prompt")
+    print()
+    
+    system_msg = args.system
+    messages = []
+    if system_msg:
+        messages.append({"role": "system", "content": system_msg})
+    
+    while True:
+        try:
+            user_input = input("You: ").strip()
+        except (KeyboardInterrupt, EOFError):
+            print("\nBye! 👋")
+            break
+        
+        if not user_input:
+            continue
+        
+        if user_input.lower() == "quit":
+            print("Bye! 👋")
+            break
+        
+        if user_input.lower() == "/reset":
+            messages = []
+            if system_msg:
+                messages.append({"role": "system", "content": system_msg})
+            model.reset_cache()
+            print("🔄 Conversation reset.\n")
+            continue
+        
+        if user_input.lower().startswith("/system "):
+            system_msg = user_input[8:].strip()
+            messages = [{"role": "system", "content": system_msg}]
+            model.reset_cache()
+            print(f"📝 System prompt updated: {system_msg}\n")
+            continue
+        
+        messages.append({"role": "user", "content": user_input})
+        prompt = format_chat(messages, template=template, add_generation_prompt=True)
+        
+        # Reset cache and regenerate from full conversation
+        model.reset_cache()
+        
+        print("Assistant: ", end="", flush=True)
+        
+        collected = []
+        start = time.time()
+        
+        def on_token(token_str):
+            print(token_str, end="", flush=True)
+            collected.append(token_str)
+        
+        model.generate_with_callback(
+            prompt,
+            on_token,
+            max_tokens=args.max_tokens,
+            temperature=args.temperature,
+            top_k=args.top_k,
+            top_p=args.top_p,
+            repetition_penalty=args.repetition_penalty,
+        )
+        
+        elapsed = time.time() - start
+        speed = len(collected) / elapsed if elapsed > 0 else 0
+        print(f"\n  [{len(collected)} tok, {speed:.1f} tok/s]\n")
+        
+        # Add assistant response to history
+        assistant_response = "".join(collected)
+        messages.append({"role": "assistant", "content": assistant_response})
+
+
+def cmd_serve(args):
+    """Start OpenAI-compatible API server."""
+    from cortex_rust.server import serve
+    serve(
+        model_path=args.model,
+        host=args.host,
+        port=args.port,
+        device=args.device,
+        tokenizer=args.tokenizer,
+        lora=getattr(args, 'lora', None),
+        q8_cache=getattr(args, 'q8_cache', False),
+        ttt=getattr(args, 'ttt', False),
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="bit-ttt",
+        description="Bit-TTT-Engine — Fast LLM inference with 1.58-bit quantization",
+    )
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+    
+    # Common arguments
+    def add_common_args(p):
+        p.add_argument("model", help="Path to GGUF model file")
+        p.add_argument("--device", default="cuda", choices=["cuda", "cpu"], help="Device (default: cuda)")
+        p.add_argument("--tokenizer", default=None, help="Path to tokenizer.json (auto-detected if not set)")
+    
+    def add_generation_args(p):
+        p.add_argument("--max-tokens", type=int, default=512, help="Max tokens to generate (default: 512)")
+        p.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature (default: 0.7)")
+        p.add_argument("--top-k", type=int, default=40, help="Top-K sampling (default: 40)")
+        p.add_argument("--top-p", type=float, default=0.9, help="Top-P (nucleus) sampling (default: 0.9)")
+        p.add_argument("--repetition-penalty", type=float, default=1.1, help="Repetition penalty (default: 1.1)")
+        p.add_argument("--template", default=None, help="Chat template (llama3/llama2/gemma2/chatml, auto-detected)")
+        p.add_argument("--system", default=None, help="System prompt")
+    
+    def add_feature_args(p):
+        p.add_argument("--lora", default=None, metavar="PATH", help="Load LoRA adapter from file")
+        p.add_argument("--q8-cache", action="store_true", help="Enable Q8 KV cache (saves ~82%% VRAM)")
+        p.add_argument("--ttt", action="store_true", help="Enable Test-Time Training (online learning)")
+    
+    # info
+    p_info = subparsers.add_parser("info", help="Show model information")
+    add_common_args(p_info)
+    p_info.set_defaults(func=cmd_info)
+    
+    # generate
+    p_gen = subparsers.add_parser("generate", aliases=["gen"], help="Generate text from a prompt")
+    add_common_args(p_gen)
+    add_generation_args(p_gen)
+    add_feature_args(p_gen)
+    p_gen.add_argument("--prompt", "-p", required=True, help="Input prompt")
+    p_gen.add_argument("--raw", action="store_true", help="Don't apply chat template")
+    p_gen.set_defaults(func=cmd_generate)
+    
+    # chat
+    p_chat = subparsers.add_parser("chat", help="Interactive chat mode")
+    add_common_args(p_chat)
+    add_generation_args(p_chat)
+    add_feature_args(p_chat)
+    p_chat.set_defaults(func=cmd_chat)
+    
+    # serve
+    p_serve = subparsers.add_parser("serve", help="Start OpenAI-compatible API server")
+    add_common_args(p_serve)
+    add_feature_args(p_serve)
+    p_serve.add_argument("--host", default="0.0.0.0", help="Bind host (default: 0.0.0.0)")
+    p_serve.add_argument("--port", type=int, default=8000, help="Bind port (default: 8000)")
+    p_serve.set_defaults(func=cmd_serve)
+    
+    args = parser.parse_args()
+    
+    if not args.command:
+        parser.print_help()
+        sys.exit(1)
+    
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

cortex_rust/engine.py

@@ -0,0 +1,253 @@
+"""High-level Python SDK for Bit-TTT-Engine.
+
+Simple, Pythonic interface for LLM inference.
+
+Quick start:
+    from cortex_rust.engine import load
+
+    model = load("model.gguf")                        # local file
+    model = load("bartowski/Llama-3-8B-Instruct-GGUF") # HuggingFace auto-download
+
+    # Generate text
+    text = model.generate("Once upon a time")
+
+    # Chat (auto-applies chat template)
+    response = model.chat([
+        {"role": "user", "content": "What is 2+2?"}
+    ])
+
+    # Stream tokens
+    for token in model.stream("Tell me a story"):
+        print(token, end="", flush=True)
+
+    # Chat stream
+    for token in model.chat_stream([
+        {"role": "system", "content": "You are a pirate."},
+        {"role": "user", "content": "Hello!"},
+    ]):
+        print(token, end="", flush=True)
+"""
+
+import os
+from typing import Dict, Iterator, List, Optional, Union
+
+Messages = List[Dict[str, str]]
+
+
+class Model:
+    """High-level model wrapper with chat, generate, and stream support."""
+
+    def __init__(self, inner, template: str, model_path: str,
+                 tokenizer=None):
+        self._inner = inner
+        self._template = template
+        self._model_path = model_path
+        self._tokenizer = tokenizer  # tokenizers.Tokenizer instance
+
+    @property
+    def config(self):
+        """Access model configuration."""
+        return self._inner.config
+
+    @property
+    def template(self) -> str:
+        """Chat template name (llama3, llama2, gemma2, chatml)."""
+        return self._template
+
+    # ---- Generation ----
+
+    def generate(self, prompt: str, max_tokens: int = 256,
+                 temperature: float = 0.7, top_p: float = 0.9,
+                 repetition_penalty: float = 1.1, **kwargs) -> str:
+        """Generate text from a prompt (no chat template applied).
+
+        Args:
+            prompt: Raw text prompt.
+            max_tokens: Maximum tokens to generate.
+            temperature: Sampling temperature (0 = greedy).
+            top_p: Nucleus sampling threshold.
+            repetition_penalty: Repetition penalty.
+
+        Returns:
+            Generated text string.
+        """
+        self._inner.reset_cache()
+        return self._inner.generate(
+            prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+        )
+
+    def chat(self, messages: Messages, max_tokens: int = 256,
+             temperature: float = 0.7, top_p: float = 0.9,
+             repetition_penalty: float = 1.1, **kwargs) -> str:
+        """Chat completion with automatic template formatting.
+
+        Args:
+            messages: List of {"role": ..., "content": ...} dicts.
+            max_tokens: Maximum tokens to generate.
+
+        Returns:
+            Assistant response text.
+        """
+        from cortex_rust.chat import format_chat
+        prompt = format_chat(messages, template=self._template)
+        return self.generate(prompt, max_tokens=max_tokens,
+                             temperature=temperature, top_p=top_p,
+                             repetition_penalty=repetition_penalty)
+
+    # ---- Streaming ----
+
+    def stream(self, prompt: str, max_tokens: int = 256,
+               temperature: float = 0.7, top_p: float = 0.9,
+               repetition_penalty: float = 1.1) -> Iterator[str]:
+        """Stream tokens from a prompt.
+
+        Uses generate_from_tokens internally and decodes token-by-token.
+        If tokenizer is not available, falls back to yielding the full output.
+
+        Yields:
+            Individual token strings as they are generated.
+        """
+        self._inner.reset_cache()
+
+        # Try token-by-token via generate_from_tokens
+        if self._tokenizer is not None:
+            encoding = self._tokenizer.encode(prompt, add_special_tokens=True)
+            prompt_ids = [int(t) for t in encoding.ids]
+            if not prompt_ids:
+                prompt_ids = [1]
+
+            gen_ids = self._inner.generate_from_tokens(
+                prompt_ids,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+            )
+            # Decode incrementally to preserve whitespace
+            all_ids = []
+            prev_text = ""
+            for tid in gen_ids:
+                all_ids.append(int(tid))
+                full_text = self._tokenizer.decode(all_ids, skip_special_tokens=True)
+                delta = full_text[len(prev_text):]
+                prev_text = full_text
+                if delta:
+                    yield delta
+        else:
+            # Fallback: yield full output at once
+            text = self._inner.generate(
+                prompt,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+            )
+            yield text
+
+    def chat_stream(self, messages: Messages, max_tokens: int = 256,
+                    temperature: float = 0.7, top_p: float = 0.9,
+                    repetition_penalty: float = 1.1) -> Iterator[str]:
+        """Stream chat completion tokens.
+
+        Yields:
+            Individual token strings.
+
+        Example:
+            for token in model.chat_stream([{"role": "user", "content": "Hi!"}]):
+                print(token, end="", flush=True)
+        """
+        from cortex_rust.chat import format_chat
+        prompt = format_chat(messages, template=self._template)
+        yield from self.stream(prompt, max_tokens=max_tokens,
+                               temperature=temperature, top_p=top_p,
+                               repetition_penalty=repetition_penalty)
+
+    # ---- Configuration ----
+
+    def enable_q8_cache(self, enabled: bool = True):
+        """Enable/disable Q8 KV cache (saves ~82% VRAM)."""
+        self._inner.enable_q8_kv_cache(enabled)
+        return self
+
+    def enable_ttt(self, enabled: bool = True, lr: float = 0.01):
+        """Enable/disable Test-Time Training."""
+        self._inner.enable_ttt(enabled)
+        return self
+
+    def create_lora(self, rank: int = 8, alpha: float = 16.0,
+                    target_modules: Optional[List[str]] = None):
+        """Create LoRA adapters for fine-tuning.
+
+        Args:
+            rank: LoRA rank (lower = faster, less capacity).
+            alpha: LoRA alpha scaling factor.
+            target_modules: List of modules to apply LoRA to.
+                Default: ["q_proj", "v_proj"]
+        """
+        modules = target_modules or ["q_proj", "v_proj"]
+        self._inner.create_lora(rank=rank, alpha=alpha, target_modules=modules)
+        return self
+
+    def reset(self):
+        """Reset KV cache and TTT state."""
+        self._inner.reset_cache()
+        return self
+
+    def __repr__(self):
+        name = os.path.basename(self._model_path)
+        return f"Model({name}, template={self._template})"
+
+
+def load(model_path: str, device: str = "cuda",
+         tokenizer: Optional[str] = None) -> Model:
+    """Load a model from a local file or HuggingFace repo.
+
+    Args:
+        model_path: Local GGUF file path or HuggingFace repo ID.
+            Examples:
+                "model.gguf"
+                "bartowski/Llama-3-8B-Instruct-GGUF"
+                "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF:Q4_K_M"
+        device: "cuda" or "cpu".
+        tokenizer: Optional path to tokenizer.json.
+
+    Returns:
+        Model instance ready for generation.
+    """
+    from cortex_rust import QGgufModel
+    from cortex_rust.chat import detect_template
+    from cortex_rust.cli import _resolve_model_path
+
+    # Resolve HF repo → local path
+    resolved = _resolve_model_path(model_path)
+
+    # Auto-detect tokenizer
+    tok = tokenizer
+    if tok is None:
+        model_dir = os.path.dirname(os.path.abspath(resolved))
+        candidate = os.path.join(model_dir, "tokenizer.json")
+        if os.path.exists(candidate):
+            tok = candidate
+
+    kwargs = {"device": device}
+    if tok:
+        kwargs["tokenizer"] = tok
+
+    inner = QGgufModel(resolved, **kwargs)
+    template = detect_template(resolved)
+
+    # Load tokenizer instance for streaming
+    tok_instance = None
+    if tok:
+        try:
+            from tokenizers import Tokenizer
+            tok_instance = Tokenizer.from_file(tok)
+        except Exception:
+            pass  # streaming will fall back to non-token mode
+
+    return Model(inner, template=template, model_path=resolved,
+                 tokenizer=tok_instance)