build-small-hackathon/microfactory-lab

1"""ZeroGPU LoRA inference backend — loads fine-tuned adapters on the Space.2 3Extends llm_zerogpu.py to wrap the base model with a PeftModel (LoRA adapter)4after loading. The adapter is only 35MB — loads in ~2 seconds after the base5model is in memory.6 7Activation: Set CHIEF_ENGINEER_LORA_REPO to a HF Hub adapter repo id.8  CHIEF_ENGINEER_LORA_REPO=kylebrodeur/microfactory-node-lora-v29 10This module is import-guarded like llm_zerogpu.py — absent deps → safe no-op.11"""12 13from __future__ import annotations14 15import json16import os17import re18 19HF_MODEL = os.environ.get("CHIEF_ENGINEER_HF_MODEL", "google/gemma-4-E4B-it")20LORA_REPO = os.environ.get("CHIEF_ENGINEER_LORA_REPO", "")21_GPU_SECONDS = int(os.environ.get("CHIEF_ENGINEER_GPU_SECONDS", "90"))22_MAX_NEW = int(os.environ.get("CHIEF_ENGINEER_MAX_NEW_TOKENS", "512"))23 24try:25    import torch  # type: ignore26    from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore27    _HAVE_HF = True28except Exception:29    torch = None  # type: ignore30    _HAVE_HF = False31 32try:33    import spaces  # type: ignore34    _HAVE_SPACES = True35except Exception:36    _HAVE_SPACES = False37 38 39def _gpu(fn):40    if _HAVE_SPACES:41        return spaces.GPU(duration=_GPU_SECONDS)(fn)42    return fn43 44 45_tok = None46_model = None47 48 49def _ensure_loaded() -> bool:50    global _tok, _model51    if not _HAVE_HF:52        return False53    if _model is not None:54        return True55    try:56        _tok = AutoTokenizer.from_pretrained(HF_MODEL)57        base = AutoModelForCausalLM.from_pretrained(58            HF_MODEL,59            dtype=getattr(torch, "bfloat16", None),60            low_cpu_mem_usage=True,61        )62        if LORA_REPO:63            from peft import PeftModel64            _model = PeftModel.from_pretrained(base, LORA_REPO)65        else:66            _model = base67        if torch is not None and torch.cuda.is_available():68            _model = _model.to("cuda")69        return True70    except Exception:71        _tok = _model = None72        return False73 74 75def is_available() -> bool:76    return _HAVE_HF77 78 79def backend_status() -> str:80    where = "ZeroGPU" if _HAVE_SPACES else "local GPU/CPU"81    if not _HAVE_HF:82        return "offline fallback · transformers/torch absent (deterministic)"83    lora_tag = f" + LoRA({LORA_REPO.split('/')[-1]})" if LORA_REPO else ""84    loaded = " (loaded)" if _model is not None else " (loads on first analyze)"85    return f"live · {HF_MODEL}{lora_tag} (transformers on {where}){loaded}"86 87 88def _build_prompt(system: str, user: str) -> str:89    messages = [{"role": "user", "content": f"{system}\n\n{user}"}]90    return _tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)91 92 93@_gpu94def _generate(system: str, user: str, temperature: float) -> str | None:95    if not _ensure_loaded():96        return None97    prompt = _build_prompt(system, user)98    if torch is not None and torch.cuda.is_available() and _model.device.type != "cuda":99        _model.to("cuda")100    inputs = _tok(prompt, return_tensors="pt").to(_model.device)101    out = _model.generate(102        **inputs,103        max_new_tokens=_MAX_NEW,104        do_sample=temperature > 0,105        temperature=max(temperature, 1e-4),106    )107    text = _tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)108    return text109 110 111@_gpu112def warm() -> str:113    if not _ensure_loaded():114        return backend_status()115    try:116        if torch is not None and torch.cuda.is_available() and _model.device.type != "cuda":117            _model.to("cuda")118        inputs = _tok("ok", return_tensors="pt").to(_model.device)119        _model.generate(**inputs, max_new_tokens=1, do_sample=False)120    except Exception:121        pass122    return backend_status()123 124 125_JSON = re.compile(r"\{.*\}", re.DOTALL)126 127 128def chat_json(system: str, user: str, temperature: float = 0.4) -> dict | None:129    try:130        text = _generate(system, user, temperature)131    except Exception:132        return None133    if not text:134        return None135    text = text.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()136    m = _JSON.search(text)137    if not m:138        return None139    try:140        return json.loads(m.group(0))141    except Exception:142        return None143