mlxsmith 0.1.0__py3-none-any.whl

mlxsmith/infer.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional
+
+from .config import ProjectConfig
+from .llm.registry import get_llm_backend
+from .models import resolve_model_spec
+
+
+@dataclass
+class ChatMessage:
+    role: str
+    content: str
+
+
+def _messages_to_prompt(messages: List[ChatMessage], tokenizer, *, use_chat_template: bool) -> str:
+    if use_chat_template and hasattr(tokenizer, "apply_chat_template"):
+        payload = [{"role": m.role, "content": m.content} for m in messages]
+        try:
+            return tokenizer.apply_chat_template(payload, tokenize=False, add_generation_prompt=True)
+        except Exception:
+            pass
+    joined = "\n".join([f"{m.role}: {m.content}" for m in messages])
+    return f"{joined}\nassistant:"
+
+
+def _load_backend(cfg: ProjectConfig, model_spec: str):
+    llm = get_llm_backend(cfg.model.backend)
+    base_model, adapter_path, _meta = resolve_model_spec(Path.cwd(), model_spec, cfg)
+    llm.load(
+        base_model,
+        max_seq_len=cfg.model.max_seq_len,
+        dtype=cfg.model.dtype,
+        trust_remote_code=cfg.model.trust_remote_code,
+    )
+    if adapter_path:
+        llm.apply_adapter(str(adapter_path))
+    return llm, base_model
+
+
+def run_prompt(
+    cfg: ProjectConfig,
+    model_spec: str,
+    prompt: str,
+    *,
+    max_new_tokens: Optional[int] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    seed: Optional[int] = None,
+) -> str:
+    llm, _base_model = _load_backend(cfg, model_spec)
+    gen = llm.generate(
+        prompt,
+        max_new_tokens=max_new_tokens or cfg.infer.max_new_tokens,
+        temperature=temperature if temperature is not None else cfg.infer.temperature,
+        top_p=top_p if top_p is not None else cfg.infer.top_p,
+        top_k=top_k if top_k is not None else cfg.infer.top_k,
+        seed=seed,
+    )
+    if gen.text.startswith(prompt):
+        return gen.text[len(prompt) :]
+    return gen.text
+
+
+def run_chat(
+    cfg: ProjectConfig,
+    model_spec: str,
+    messages: List[ChatMessage],
+    *,
+    max_new_tokens: Optional[int] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    seed: Optional[int] = None,
+) -> str:
+    llm, _base_model = _load_backend(cfg, model_spec)
+    prompt = _messages_to_prompt(messages, llm.tokenizer, use_chat_template=cfg.model.use_chat_template)
+    gen = llm.generate(
+        prompt,
+        max_new_tokens=max_new_tokens or cfg.infer.max_new_tokens,
+        temperature=temperature if temperature is not None else cfg.infer.temperature,
+        top_p=top_p if top_p is not None else cfg.infer.top_p,
+        top_k=top_k if top_k is not None else cfg.infer.top_k,
+        seed=seed,
+    )
+    if gen.text.startswith(prompt):
+        return gen.text[len(prompt) :]
+    return gen.text
+
+
+def chat_repl(
+    cfg: ProjectConfig,
+    model_spec: str,
+    *,
+    system: Optional[str] = None,
+    max_new_tokens: Optional[int] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    seed: Optional[int] = None,
+    max_turns: Optional[int] = None,
+) -> None:
+    llm, _base_model = _load_backend(cfg, model_spec)
+    messages: List[ChatMessage] = []
+    if system:
+        messages.append(ChatMessage(role="system", content=system))
+
+    turns = 0
+    while True:
+        if max_turns is not None and turns >= max_turns:
+            break
+        try:
+            user = input("user> ").strip()
+        except EOFError:
+            break
+        if not user:
+            continue
+        if user.lower() in {"/exit", "/quit", "exit", "quit"}:
+            break
+        messages.append(ChatMessage(role="user", content=user))
+        prompt = _messages_to_prompt(messages, llm.tokenizer, use_chat_template=cfg.model.use_chat_template)
+        gen = llm.generate(
+            prompt,
+            max_new_tokens=max_new_tokens or cfg.infer.max_new_tokens,
+            temperature=temperature if temperature is not None else cfg.infer.temperature,
+            top_p=top_p if top_p is not None else cfg.infer.top_p,
+            top_k=top_k if top_k is not None else cfg.infer.top_k,
+            seed=seed,
+        )
+        if gen.text.startswith(prompt):
+            reply = gen.text[len(prompt) :]
+        else:
+            reply = gen.text
+        reply = reply.strip()
+        print(f"assistant> {reply}\n")
+        messages.append(ChatMessage(role="assistant", content=reply))
+        turns += 1

mlxsmith/llm/__init__.py

@@ -0,0 +1,16 @@
+"""Model backends for mlxsmith."""
+
+from .backend import LLMBackend, Generation, BackendNotAvailable, DecodingConfig
+from .mlx_lm_backend import MlxLMBackend
+from .mock_backend import MockBackend
+from .registry import get_llm_backend
+
+__all__ = [
+    "LLMBackend",
+    "Generation",
+    "BackendNotAvailable",
+    "MlxLMBackend",
+    "MockBackend",
+    "DecodingConfig",
+    "get_llm_backend",
+]

mlxsmith/llm/backend.py

@@ -0,0 +1,126 @@
+"""LLM backend abstraction.
+
+We keep this intentionally small so mlxsmith can support multiple model loaders
+without hard-binding to one ecosystem.
+
+Primary target: mlx-lm (HF -> MLX format + common chat models).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol, Sequence, Any, Optional, List, Dict
+
+
+@dataclass
+class Generation:
+    text: str
+    token_ids: list[int]
+    prompt_len: int
+    # Log-probabilities for generated tokens only (len = completion tokens), if available.
+    logprobs: list[float] | None = None
+    # Top-k logprobs per token: list of {token_str: logprob} dicts for each generated token.
+    top_k_logprobs: List[Dict[str, float]] | None = None
+
+
+@dataclass
+class DecodingConfig:
+    max_new_tokens: int = 256
+    temperature: float = 0.8
+    top_p: float = 1.0
+    top_k: Optional[int] = None
+    seed: Optional[int] = None
+    stop: Optional[Sequence[str]] = None
+
+
+class LLMBackend(Protocol):
+    name: str
+
+    def load(self, model_id_or_path: str, *, max_seq_len: int | None = None, dtype: str | None = None) -> None:
+        """Load model + tokenizer into memory."""
+
+    def encode(self, text: str) -> list[int]:
+        """Tokenize text -> token ids."""
+
+    def decode(self, ids: Sequence[int]) -> str:
+        """Token ids -> text."""
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        max_new_tokens: int = 256,
+        temperature: float = 0.8,
+        top_p: float = 1.0,
+        top_k: int | None = None,
+        seed: int | None = None,
+    ) -> Generation:
+        """Sample a completion and return ids for prompt+completion."""
+
+    def generate_with_logprobs(
+        self,
+        prompt: str,
+        *,
+        max_new_tokens: int = 256,
+        temperature: float = 0.8,
+        top_p: float = 1.0,
+        top_k: int | None = None,
+        seed: int | None = None,
+        logprobs: int = 0,  # Number of top logprobs to return per token (0 = just the sampled token)
+    ) -> Generation:
+        """Sample a completion and include per-token logprobs when available.
+        
+        Args:
+            prompt: Input prompt text
+            max_new_tokens: Maximum number of tokens to generate
+            temperature: Sampling temperature
+            top_p: Nucleus sampling parameter
+            top_k: Top-k sampling parameter
+            seed: Random seed for reproducibility
+            logprobs: Number of top logprobs to return per token (0 = return only sampled token's logprob)
+            
+        Returns:
+            Generation with logprobs and optionally top_k_logprobs
+        """
+
+    def sft_loss(self, token_ids: Sequence[int], *, train_on_prompt: bool, prompt_len: int) -> Any:
+        """Return a scalar loss suitable for backprop (MLX array)."""
+
+    def rl_loss(self, token_ids: Sequence[int], *, prompt_len: int, advantage: float) -> Any:
+        """Return policy-gradient-style loss for a sampled trajectory."""
+
+    def sequence_logprob(self, token_ids: Sequence[int], *, prompt_len: int) -> Any:
+        """Return log-probability sum of response tokens (differentiable)."""
+
+    def token_logprobs(
+        self,
+        token_ids: Sequence[int],
+        *,
+        prompt_len: int,
+        top_k: int = 0,
+        include_prompt: bool = False,
+    ) -> tuple[list[float], List[Dict[str, float]] | None]:
+        """Return per-token logprobs (and optional top-k logprobs).
+
+        Args:
+            token_ids: Full token sequence (prompt + completion).
+            prompt_len: Prompt length in tokens.
+            top_k: Number of top logprobs per token to return (0 = none).
+            include_prompt: If True, include prompt tokens; otherwise only response tokens.
+        """
+
+    def value_and_grad(self, loss_fn) -> tuple[Any, Any | None]:
+        """Return (loss, grads) using backend autograd when available."""
+
+    def optimizer_and_params(self, *, lr: float, weight_decay: float = 0.0) -> tuple[Any, Any]:
+        """Return (optimizer, trainable_params_tree)."""
+
+    def apply_grads(self, optimizer: Any, grads: Any) -> None:
+        """Update model parameters given gradients."""
+
+    def save_adapter(self, out_dir: str, *, metadata: dict | None = None) -> None:
+        """Persist adapter weights (LoRA) to out_dir."""
+
+
+class BackendNotAvailable(RuntimeError):
+    pass

mlxsmith/llm/interface.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, List, Optional, Dict
+
+from .registry import get_llm_backend
+from .backend import DecodingConfig, Generation
+from ..config import ProjectConfig
+from ..models import resolve_model_spec
+from ..train.lora import load_adapter_config
+
+
+@dataclass
+class LoadedModel:
+    backend: any
+    base_model: str
+    adapter_path: Optional[str]
+
+
+def load_base_model(model_id_or_path: str, cfg: ProjectConfig) -> LoadedModel:
+    """Load a base model (and optionally adapter) with the configured backend."""
+    backend = get_llm_backend(cfg.model.backend)
+    from pathlib import Path
+
+    base_model, adapter_path, _meta = resolve_model_spec(Path.cwd(), model_id_or_path, cfg)
+    backend.load(
+        base_model,
+        max_seq_len=cfg.model.max_seq_len,
+        dtype=cfg.model.dtype,
+        trust_remote_code=cfg.model.trust_remote_code,
+    )
+    return LoadedModel(backend=backend, base_model=base_model, adapter_path=str(adapter_path) if adapter_path else None)
+
+
+def load_adapter(adapter_path: str):
+    """Load adapter config without applying it."""
+    return load_adapter_config(adapter_path)
+
+
+def apply_adapter(backend, adapter_path: str) -> None:
+    backend.apply_adapter(adapter_path)
+
+
+def generate(
+    backend, 
+    tokenizer, 
+    prompts: Iterable[str], 
+    decoding_config: DecodingConfig,
+    logprobs: int = 0,
+) -> List[Generation]:
+    """Generate completions for prompts.
+    
+    Args:
+        backend: The LLM backend
+        tokenizer: Tokenizer instance
+        prompts: Iterable of prompt strings
+        decoding_config: Decoding configuration
+        logprobs: Number of top logprobs to return per token (0 = none)
+        
+    Returns:
+        List of Generation results
+    """
+    results = []
+    for p in prompts:
+        if logprobs > 0:
+            results.append(
+                backend.generate_with_logprobs(
+                    p,
+                    max_new_tokens=decoding_config.max_new_tokens,
+                    temperature=decoding_config.temperature,
+                    top_p=decoding_config.top_p,
+                    top_k_sampling=decoding_config.top_k,
+                    seed=decoding_config.seed,
+                    logprobs=logprobs,
+                )
+            )
+        else:
+            results.append(
+                backend.generate(
+                    p,
+                    max_new_tokens=decoding_config.max_new_tokens,
+                    temperature=decoding_config.temperature,
+                    top_p=decoding_config.top_p,
+                    top_k=decoding_config.top_k,
+                    seed=decoding_config.seed,
+                )
+            )
+    return results
+
+
+def chat(
+    backend,
+    messages: List[Dict[str, str]],
+    decoding_config: DecodingConfig,
+    logprobs: int = 0,
+    use_chat_template: bool = True,
+) -> Generation:
+    """Generate a chat completion.
+    
+    Args:
+        backend: The LLM backend
+        messages: List of message dicts with 'role' and 'content' keys
+        decoding_config: Decoding configuration
+        logprobs: Number of top logprobs to return per token (0 = none)
+        use_chat_template: Whether to use the model's chat template
+        
+    Returns:
+        Generation result
+    """
+    # Convert messages to prompt
+    prompt = _messages_to_prompt(backend.tokenizer, messages, use_chat_template=use_chat_template)
+    
+    if logprobs > 0:
+        return backend.generate_with_logprobs(
+            prompt,
+            max_new_tokens=decoding_config.max_new_tokens,
+            temperature=decoding_config.temperature,
+            top_p=decoding_config.top_p,
+            top_k_sampling=decoding_config.top_k,
+            seed=decoding_config.seed,
+            logprobs=logprobs,
+        )
+    else:
+        return backend.generate(
+            prompt,
+            max_new_tokens=decoding_config.max_new_tokens,
+            temperature=decoding_config.temperature,
+            top_p=decoding_config.top_p,
+            top_k=decoding_config.top_k,
+            seed=decoding_config.seed,
+        )
+
+
+def _messages_to_prompt(
+    tokenizer,
+    messages: List[Dict[str, str]],
+    *,
+    use_chat_template: bool = True
+) -> str:
+    """Convert chat messages to prompt string."""
+    if use_chat_template and hasattr(tokenizer, "apply_chat_template"):
+        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    # Fallback
+    return "\n".join([f"{m['role']}: {m['content']}" for m in messages]) + "\nassistant:"
+
+
+@dataclass
+class LogprobResult:
+    """Result from logprobs computation."""
+    token_logprobs: List[float]
+    top_k_logprobs: Optional[List[Dict[str, float]]] = None
+    text: Optional[str] = None
+
+
+def compute_logprobs(
+    backend,
+    prompt: str,
+    completion: str,
+    top_k: int = 0,
+    max_seq_len: Optional[int] = None,
+) -> LogprobResult:
+    """Compute logprobs for a prompt-completion pair.
+    
+    Args:
+        backend: The LLM backend
+        prompt: Prompt text
+        completion: Completion text
+        top_k: Number of top logprobs per token to return (0 = none)
+        max_seq_len: Maximum sequence length
+        
+    Returns:
+        LogprobResult with token logprobs and optionally top-k logprobs
+    """
+    prompt_ids = backend.encode(prompt)
+    ids = backend.encode(prompt + completion)
+    
+    # Truncate if needed
+    if max_seq_len and len(ids) > max_seq_len:
+        overflow = len(ids) - max_seq_len
+        ids = ids[overflow:]
+        prompt_len = max(0, len(prompt_ids) - overflow)
+    else:
+        prompt_len = len(prompt_ids)
+    
+    # Get generation with logprobs
+    full_text = backend.decode(ids)
+    
+    # Use backend's sequence_logprob if available
+    seq_logprob = backend.sequence_logprob(ids, prompt_len=prompt_len)
+    
+    # For per-token logprobs, we'd need to do a forward pass
+    # This is a simplified version
+    token_logprobs = []
+    if hasattr(backend, '_response_logprobs'):
+        token_logprobs = backend._response_logprobs(ids, prompt_len=prompt_len)
+    
+    # Get top-k logprobs if requested
+    top_k_logprobs = None
+    if top_k > 0 and hasattr(backend, 'generate_with_logprobs'):
+        # Generate to get top-k for each position
+        gen = backend.generate_with_logprobs(
+            prompt,
+            max_new_tokens=len(ids) - prompt_len,
+            logprobs=top_k,
+        )
+        top_k_logprobs = gen.top_k_logprobs
+    
+    return LogprobResult(
+        token_logprobs=token_logprobs,
+        top_k_logprobs=top_k_logprobs,
+        text=completion,
+    )