data-morph-gemma 0.1.0__py3-none-any.whl

datamorph/models/gemma_mlx.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import os
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+_LOCAL_PATH = Path(__file__).resolve().parents[2] / "models" / "gemma-4-e2b-it-bf16"
+# Base model dir; override via GEMMA_MLX_MODEL to point at e.g. a quantized build.
+MODEL_ID = os.environ.get("GEMMA_MLX_MODEL", str(_LOCAL_PATH))
+
+# When the model is a stripped text-only build (gemma4_text), load it through
+# mlx_lm instead of mlx_vlm — there is no vision/audio tower to wrap. Set
+# GEMMA_TEXT_ONLY=1 to select this path (used by the W7 text-only student).
+TEXT_ONLY = os.environ.get("GEMMA_TEXT_ONLY", "").lower() in ("1", "true", "yes")
+
+# Model selection lives in _state so it can be changed at runtime (via use_model)
+# without re-importing the module. The env vars seed the defaults, so existing
+# scripts/notebooks that set GEMMA_MLX_MODEL / GEMMA_TEXT_ONLY keep working.
+_state: dict = {
+    "model": None,
+    "processor": None,
+    "load_sec": None,
+    "adapter": None,
+    "model_id": MODEL_ID,
+    "text_only": TEXT_ONLY,
+}
+
+
+def use_adapter(adapter_path: str | None) -> None:
+    """Select a LoRA adapter directory for inference (None = base model).
+
+    The directory must contain ``adapters.safetensors`` + ``adapter_config.json``.
+    Changing the selection forces the model to reload on the next ``generate`` call.
+    """
+    if adapter_path != _state["adapter"]:
+        _state["adapter"] = adapter_path
+        _state["model"] = None  # force reload with the new adapter (or base)
+
+
+def use_model(model_id: str, *, text_only: bool = True) -> None:
+    """Select which model to load (local dir path or HF id).
+
+    Mirrors ``use_adapter``: changing the selection forces a reload on the next
+    ``generate`` call. ``text_only=True`` loads via mlx_lm (the stripped W7
+    student); ``False`` uses mlx_vlm (the multimodal base).
+    """
+    if model_id != _state["model_id"] or text_only != _state["text_only"]:
+        _state["model_id"] = model_id
+        _state["text_only"] = text_only
+        _state["model"] = None  # force reload with the new model
+
+
+@dataclass
+class GenerationResult:
+    text: str
+    n_prompt_tokens: int
+    n_generated_tokens: int
+    elapsed_sec: float
+    tokens_per_sec: float
+    model_id: str
+    truncated: bool
+
+
+def _ensure_loaded() -> None:
+    if _state["model"] is not None:
+        return
+    t0 = time.time()
+    model_id = _state["model_id"]
+    text_only = _state["text_only"]
+    if text_only:
+        import mlx_lm  # lazy import
+
+        loaded = mlx_lm.load(model_id, adapter_path=_state["adapter"])
+        model, processor = loaded[0], loaded[1]
+    else:
+        from mlx_vlm import load  # lazy import
+
+        model, processor = load(model_id, adapter_path=_state["adapter"])
+    _state["model"] = model
+    _state["processor"] = processor
+    _state["load_sec"] = round(time.time() - t0, 2)
+    tag = f" + adapter {_state['adapter']}" if _state["adapter"] else ""
+    backend = "mlx_lm/text" if text_only else "mlx_vlm"
+    print(f"[gemma_mlx] loaded {model_id}{tag} via {backend} in {_state['load_sec']}s")
+
+
+def _generate_text_only(messages: list[dict], max_tokens: int) -> GenerationResult:
+    """Greedy generation via mlx_lm on the stripped text-only student."""
+    import mlx_lm  # lazy
+    from mlx_lm.sample_utils import make_sampler  # lazy
+
+    model = _state["model"]
+    tok = _state["processor"]
+    prompt = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+
+    t0 = time.time()
+    text = mlx_lm.generate(
+        model,
+        tok,
+        str(prompt),
+        max_tokens=max_tokens,
+        sampler=make_sampler(temp=0.0),  # greedy, matches the temperature=0.0 VLM path
+        verbose=False,
+    )
+    elapsed = time.time() - t0
+
+    n_prompt = len(tok.encode(str(prompt)))
+    n_gen = len(tok.encode(text)) if text else 0
+    tps = round(n_gen / elapsed, 2) if elapsed > 0 else 0.0
+    return GenerationResult(
+        text=text,
+        n_prompt_tokens=n_prompt,
+        n_generated_tokens=n_gen,
+        elapsed_sec=round(elapsed, 2),
+        tokens_per_sec=tps,
+        model_id=_state["model_id"],
+        truncated=n_gen >= max_tokens,
+    )
+
+
+def generate(messages: list[dict], max_tokens: int = 4096) -> GenerationResult:
+    """Run greedy generation on the given chat messages (text-only)."""
+    _ensure_loaded()
+    if _state["text_only"]:
+        return _generate_text_only(messages, max_tokens)
+
+    from mlx_vlm import generate as vlm_generate  # lazy
+    from mlx_vlm.prompt_utils import apply_chat_template  # lazy
+
+    model = _state["model"]
+    processor = _state["processor"]
+
+    prompt = apply_chat_template(processor, model.config, messages, num_images=0)
+
+    t0 = time.time()
+    result = vlm_generate(
+        model=model,
+        processor=processor,
+        prompt=str(prompt),
+        max_tokens=max_tokens,
+        temperature=0.0,
+        verbose=False,
+    )
+    elapsed = time.time() - t0
+
+    n_prompt = int(getattr(result, "prompt_tokens", 0) or 0)
+    n_gen = int(getattr(result, "generation_tokens", 0) or 0)
+    if elapsed > 0:
+        tps = float(getattr(result, "generation_tps", n_gen / elapsed) or 0.0)
+    else:
+        tps = 0.0
+    truncated = n_gen >= max_tokens
+
+    return GenerationResult(
+        text=result.text,
+        n_prompt_tokens=n_prompt,
+        n_generated_tokens=n_gen,
+        elapsed_sec=round(elapsed, 2),
+        tokens_per_sec=round(tps, 2),
+        model_id=_state["model_id"],
+        truncated=truncated,
+    )

datamorph/models/gemma_script_teacher.py

@@ -0,0 +1,100 @@
+"""Gemma script-generation 'teacher' — the un-fine-tuned student on the new pipeline.
+
+Drop-in replacement for ``datamorph.data.teacher_script.call_script_teacher`` (same
+``ScriptResult`` contract and ``(envelope, instruction, output_format, *, feedback)``
+signature), but the script author is the local Gemma model via MLX instead of Opus.
+
+Used to establish the **zero-shot base-student baseline** on the Stage-3..5 pipeline
+(envelope -> script -> sandbox -> metrics), before any LoRA fine-tuning. The base
+model has not learned the skill, so — like the W2 Gemma baseline — the skill text is
+folded into the prompt (Gemma's chat template has no system role and it cannot Read
+files the way the Opus CLI can).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from datamorph.data.teacher_script import ScriptResult, parse_teacher_output
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+SKILL_REL_PATH = "skills/script_generation_teacher.md"
+
+_SKILL_CACHE: dict[str, str] = {}
+
+
+def _skill_text() -> str:
+    if "text" not in _SKILL_CACHE:
+        _SKILL_CACHE["text"] = (PROJECT_ROOT / SKILL_REL_PATH).read_text(encoding="utf-8")
+    return _SKILL_CACHE["text"]
+
+
+def build_gemma_prompt(
+    envelope: dict[str, Any],
+    instruction: str,
+    output_format: str,
+    feedback: str | None = None,
+) -> str:
+    """Skill text + the inference-style task (envelope + instruction + output contract).
+
+    The task half mirrors ``format_pairs.to_chat_record`` (what the fine-tuned model
+    will see at inference); the skill text is prepended so the *un-trained* base model
+    has the instructions in-context.
+    """
+    env_json = json.dumps(envelope, indent=2, default=str)
+    fb = (
+        f"\n\nYour previous attempt failed: {feedback}\nWrite a corrected response.\n"
+        if feedback
+        else ""
+    )
+    return (
+        f"{_skill_text()}\n\n---\n\n"
+        f"Metadata envelope:\n```json\n{env_json}\n```\n\n"
+        f"Task: {instruction}\n"
+        f"Target output format: {output_format.upper()}.\n"
+        f"Write a Python conversion script that reads the input file path from "
+        f"sys.argv[1] and writes the converted output to sys.argv[2], using only the "
+        f"Python standard library and pandas. Respond with exactly an "
+        f"<analysis>...</analysis> block followed by a <script>...</script> block. "
+        f"No prose outside the tags."
+        f"{fb}"
+    )
+
+
+def call_gemma_script_teacher(
+    envelope: dict[str, Any],
+    instruction: str,
+    output_format: str,
+    *,
+    feedback: str | None = None,
+    max_tokens: int = 4096,
+) -> ScriptResult:
+    """Generate <analysis> + <script> with the local Gemma model; parse into a ScriptResult."""
+    from datamorph.models.gemma_mlx import generate as mlx_generate  # lazy: keep MLX import optional
+
+    prompt = build_gemma_prompt(envelope, instruction, output_format, feedback)
+    try:
+        gen = mlx_generate([{"role": "user", "content": prompt}], max_tokens=max_tokens)
+    except Exception as e:  # MLX / generation failure -> looks like "no script"
+        return ScriptResult("", "", "", -1, f"gemma generate raised: {e!r}", {})
+
+    analysis, script = parse_teacher_output(gen.text)
+    return ScriptResult(
+        analysis=analysis,
+        script=script,
+        raw_output=gen.text,
+        returncode=0,
+        stderr="",
+        raw_payload={
+            "model_id": gen.model_id,
+            "gemma_meta": {
+                "n_prompt_tokens": gen.n_prompt_tokens,
+                "n_generated_tokens": gen.n_generated_tokens,
+                "tokens_per_sec": gen.tokens_per_sec,
+                "elapsed_sec": gen.elapsed_sec,
+                "truncated": gen.truncated,
+            },
+        },
+    )