@lzhzzzzwill/cofos 1.3.6 → 1.3.8

package/README.md

@@ -77,6 +77,14 @@ Show model reasoning scaffold when emitted:
 npx @lzhzzzzwill/cofos --show-reasoning
 ```
 
+Select the inference device explicitly:
+
+```bash
+npx @lzhzzzzwill/cofos --device cuda   # CUDA GPU machine
+npx @lzhzzzzwill/cofos --device mps    # Apple Silicon Mac
+npx @lzhzzzzwill/cofos --device cpu    # CPU fallback, slow
+```
+
 Ask once and exit, useful for scripts and smoke tests:
 
 ```bash
@@ -95,7 +103,8 @@ On startup, the CLI does the following:
 5. If `--pdf-dir` is provided, scans PDF names, sizes, and mtimes.
 6. Re-parses PDFs only when the folder contents changed or `--rebuild-pdf-index` is passed.
 7. Merges ROS-related PDF chunks into BM25 retrieval.
-8. Loads the model once and keeps the Python backend alive for the whole chat session.
+8. Selects the inference backend: CUDA first, then Apple Silicon MPS, then CPU.
+9. Loads the model once and keeps the Python backend alive for the whole chat session.
 
 PDF files are read locally and are not uploaded anywhere. Set `COFOS_CACHE_DIR` to override the COFOS cache root, or `HF_HOME` to override only the Hugging Face model cache.
 
@@ -133,6 +142,8 @@ cofos \
 | `--cache-dir` | Folder for model, KG/BM25, and parsed PDF cache; defaults to `./.cofos` |
 | `--rebuild-pdf-index` | Force PDF parsing and BM25 rebuild |
 | `--top-k` | Retrieval top-k |
+| `--device` | `auto`, `cuda`, `mps`, or `cpu`; defaults to `auto` |
+| `--dtype` | `auto`, `float32`, `float16`, or `bfloat16`; defaults to `auto` |
 | `-p`, `--prompt` | Ask one question and exit |
 | `--show-reasoning` | Show emitted reasoning scaffold in dim text before the final answer |
 | `--max-new-tokens` | Override generation length |

package/bin/cofos.js

@@ -216,7 +216,7 @@ function createPlainUi() {
 async function main() {
   if (process.argv.includes("-h") || process.argv.includes("--help")) {
     console.log("\n  COFOS CLI — 9B + RAG chat (Willlzh/COFOS)\n");
-    console.log("  Usage:  cofos [-p QUESTION] [--show-reasoning] [--model PATH] [--config PATH] [--pdf-dir DIR] [--cache-dir DIR] [--top-k N] [--no-rag]\n");
+    console.log("  Usage:  cofos [-p QUESTION] [--show-reasoning] [--device DEVICE] [--dtype DTYPE] [--model PATH] [--config PATH] [--pdf-dir DIR] [--cache-dir DIR] [--top-k N] [--no-rag]\n");
     console.log("  Commands in chat:");
     console.log("    /help        Show this help");
     console.log("    /exit        Quit");
@@ -229,6 +229,8 @@ async function main() {
     console.log("    --pdf-dir DIR          Add PDFs from DIR as extra BM25 evidence");
     console.log("    --cache-dir DIR        Store model/data cache here (default: ./.cofos)");
     console.log("    --max-new-tokens N     Override generation length");
+    console.log("    --device DEVICE        auto|cuda|mps|cpu (default: auto; CUDA first, then Mac MPS)");
+    console.log("    --dtype DTYPE          auto|float32|float16|bfloat16 (default: auto)");
     console.log("    --show-reasoning       Show model reasoning scaffold in dim text when emitted");
     console.log("    -p, --prompt TEXT      Ask once and exit");
     console.log("    --rebuild-pdf-index    Force reparsing PDFs and rebuilding BM25\n");
@@ -247,6 +249,8 @@ async function main() {
   let cacheDir = process.env.COFOS_CACHE_DIR ? resolve(process.env.COFOS_CACHE_DIR) : resolve(process.cwd(), ".cofos");
   let rebuildPdfIndex = false;
   let showReasoning = false;
+  let devicePreference = process.env.COFOS_DEVICE || "auto";
+  let dtypePreference = process.env.COFOS_DTYPE || "auto";
   let singlePrompt = null;
   const backendArgs = [chatScript];
 
@@ -280,6 +284,12 @@ async function main() {
       rebuildPdfIndex = true;
     } else if (arg === "--show-reasoning") {
       showReasoning = true;
+    } else if (arg === "--device") {
+      devicePreference = value(++i, arg).toLowerCase();
+      if (!["auto", "cuda", "mps", "cpu", "gpu", "metal"].includes(devicePreference)) die(`--device must be one of auto, cuda, mps, cpu; got "${devicePreference}"`);
+    } else if (arg === "--dtype") {
+      dtypePreference = value(++i, arg).toLowerCase();
+      if (!["auto", "fp32", "float32", "fp16", "float16", "bf16", "bfloat16"].includes(dtypePreference)) die(`--dtype must be one of auto, float32, float16, bfloat16; got "${dtypePreference}"`);
     } else if (arg === "-p" || arg === "--prompt") {
       singlePrompt = value(++i, arg);
     } else if (arg === "--no-rag") {
@@ -318,6 +328,8 @@ async function main() {
     PYTHONUNBUFFERED: process.env.PYTHONUNBUFFERED || "1",
     PYTHONIOENCODING: process.env.PYTHONIOENCODING || "utf-8",
     COFOS_SHOW_REASONING: showReasoning ? "1" : (process.env.COFOS_SHOW_REASONING || "0"),
+    COFOS_DEVICE: devicePreference,
+    COFOS_DTYPE: dtypePreference,
   };
   const proc = spawn(py, backendArgs, {
     stdio: ["pipe", "pipe", "pipe"],
@@ -327,7 +339,7 @@ async function main() {
   proc.stdout.setEncoding("utf-8");
   proc.stderr.setEncoding("utf-8");
 
-  let buf = "", ready = false, generating = false, responseOpen = false, commandOpen = false, closed = false;
+  let buf = "", ready = false, generating = false, responseOpen = false, commandOpen = false, closed = false, backendExited = false;
   let currentAnswer = "", assistantStart = -1, plainAssistantMode = "undecided", plainAssistantPending = "";
   let plainReasonBuffer = "", plainReasonPrinted = 0;
   const inputQ = [];
@@ -356,8 +368,12 @@ async function main() {
   function handleLocalOnly(rawInput) {
     const input = String(rawInput || "").trim();
     if (!input) return true;
+    if (backendExited && !/^\/(exit|quit)$/.test(input)) {
+      ui.print("• COFOS backend is not running. Restart `cofos` after fixing the error above, or run /exit.");
+      return true;
+    }
     if (input === "/") {
-      ui.print("Commands: " + SLASH_COMMANDS.map(({ cmd }) => cmd).join("  ") + "\nType a prefix such as /r or /to, then press Tab for completion.");
+      ui.print("• Type a command prefix such as /r or /to, then press Tab. Run /help for all commands.");
       return true;
     }
     if (/^\/(exit|quit)$/.test(input)) { shutdown(); return true; }
@@ -370,7 +386,7 @@ async function main() {
       return true;
     }
     if (/^\/(info)$/.test(input)) {
-      ui.print(`\n  COFOS CLI — v${pkg.version}\n  Model:  ${modelPath} (Hugging Face)\n  RAG:    Willlzh/COFOS_data/runtime (KG + BM25) + optional --pdf-dir BM25 evidence\n  Cache:  ${cacheDir}\n  HF:     ${backendEnv.HF_HOME}\n  Engine: StudentAdapterInference + transformers + PyTorch\n  Status: ${ready ? "ready" : "loading"}\n  Hist:   ~/.cofos_history.jsonl\n`);
+      ui.print(`\n  COFOS CLI — v${pkg.version}\n  Model:  ${modelPath} (Hugging Face)\n  RAG:    Willlzh/COFOS_data/runtime (KG + BM25) + optional --pdf-dir BM25 evidence\n  Cache:  ${cacheDir}\n  HF:     ${backendEnv.HF_HOME}\n  Engine: StudentAdapterInference + transformers + PyTorch\n  Device: ${devicePreference}\n  Dtype:  ${dtypePreference}\n  Status: ${ready ? "ready" : "loading"}\n  Hist:   ~/.cofos_history.jsonl\n`);
       return true;
     }
     if (input.startsWith("/") && !/^\/(clear|save|topk)(?:\s|$)/.test(input) && !/^\/rag\s+(?:on|off)$/.test(input)) {
@@ -397,8 +413,12 @@ async function main() {
     input = input.trim();
     if (!input) return false;
 
+    if (backendExited && !/^\/(exit|quit)$/.test(input)) {
+      ui.print("• COFOS backend is not running. Restart `cofos` after fixing the error above, or run /exit.");
+      return true;
+    }
     if (input === "/") {
-      ui.print("Commands: " + SLASH_COMMANDS.map(({ cmd }) => cmd).join("  ") + "\nType a prefix such as /r or /to, then press Tab for completion.");
+      ui.print("• Type a command prefix such as /r or /to, then press Tab. Run /help for all commands.");
       return false;
     }
     if (/^\/(exit|quit)$/.test(input)) { shutdown(); return true; }
@@ -439,7 +459,7 @@ async function main() {
       return true;
     }
     if (/^\/(info)$/.test(input)) {
-      ui.print(`\n  COFOS CLI — v${pkg.version}\n  Model:  ${modelPath} (Hugging Face)\n  RAG:    Willlzh/COFOS_data/runtime (KG + BM25) + optional --pdf-dir BM25 evidence\n  Cache:  ${cacheDir}\n  HF:     ${backendEnv.HF_HOME}\n  Engine: StudentAdapterInference + transformers + PyTorch\n  Hist:   ~/.cofos_history.jsonl\n`);
+      ui.print(`\n  COFOS CLI — v${pkg.version}\n  Model:  ${modelPath} (Hugging Face)\n  RAG:    Willlzh/COFOS_data/runtime (KG + BM25) + optional --pdf-dir BM25 evidence\n  Cache:  ${cacheDir}\n  HF:     ${backendEnv.HF_HOME}\n  Engine: StudentAdapterInference + transformers + PyTorch\n  Device: ${devicePreference}\n  Dtype:  ${dtypePreference}\n  Hist:   ~/.cofos_history.jsonl\n`);
       return false;
     }
     if (input.startsWith("/")) {
@@ -465,6 +485,11 @@ async function main() {
       if (!closed) ui.focus();
       return;
     }
+    if (backendExited) {
+      ui.print("• COFOS backend is not running. Restart `cofos` after fixing the error above, or run /exit.");
+      if (!closed) ui.focus();
+      return;
+    }
     if (!ready || generating) {
       inputQ.push(line);
       ui.print(ready ? "• Queued until the current answer finishes." : "• Queued until the model is ready. Use /exit to quit while loading.");
@@ -658,8 +683,17 @@ async function main() {
   proc.on("error", (e) => { ui.print("[cofos] Python error: " + e.message); shutdown(); });
   proc.on("exit", (c) => {
     if (closed) return;
+    backendExited = true;
+    ready = false;
+    generating = false;
+    inputQ.length = 0;
     ui.print(c && c !== 0 ? "[cofos] Python exited with code " + c : "[cofos] Python backend exited.");
-    ui.print("Press Ctrl+C to close.");
+    if (c && c !== 0) {
+      ui.print("[cofos] Backend failed before the model became ready. On Mac this usually means the 9B model could not fit in available memory and Transformers tried to offload the whole model to disk.");
+      ui.print("[cofos] Try a machine with more unified/GPU memory, force CPU with --device cpu for a slower attempt, or use a smaller/quantized model. Use /exit to close this session.");
+    } else {
+      ui.print("Press Ctrl+C to close.");
+    }
     ui.focus();
     process.exitCode = c ?? 0;
   });

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lzhzzzzwill/cofos",
-  "version": "1.3.6",
+  "version": "1.3.8",
   "description": "COFOS 9B + RAG chat CLI with optional local PDF ingestion.",
   "type": "module",
   "bin": {

package/runtime/src/inference/student_adapter_inference.py

@@ -225,14 +225,32 @@ class StudentAdapterInference:
             raise FileNotFoundError(f"Base model path not found: {self.model_path}")
 
     @staticmethod
-    def _select_dtype() -> torch.dtype:
-        """Pick a stable dtype for the current runtime.
+    def _select_device() -> str:
+        requested = os.environ.get("COFOS_DEVICE", "auto").strip().lower()
+        aliases = {"gpu": "cuda", "metal": "mps"}
+        requested = aliases.get(requested, requested)
+        if requested not in {"", "auto", "cuda", "mps", "cpu"}:
+            logging.warning("Ignoring unsupported COFOS_DEVICE=%s; using auto", requested)
+            requested = "auto"
+        if requested == "cuda":
+            if not torch.cuda.is_available():
+                raise RuntimeError("COFOS_DEVICE=cuda was requested, but CUDA is not available.")
+            return "cuda"
+        if requested == "mps":
+            if not getattr(torch.backends, "mps", None) or not torch.backends.mps.is_available():
+                raise RuntimeError("COFOS_DEVICE=mps was requested, but PyTorch MPS is not available.")
+            return "mps"
+        if requested == "cpu":
+            return "cpu"
+        if torch.cuda.is_available():
+            return "cuda"
+        if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+            return "mps"
+        return "cpu"
 
-        CUDA can safely use bf16/fp16.  On Mac/CPU, ``device_map=auto`` may place
-        layers on CPU or disk; fp16 there can produce NaN probabilities or
-        garbage tokens.  Prefer float32 outside CUDA unless the user explicitly
-        opts into another dtype with COFOS_DTYPE.
-        """
+    @staticmethod
+    def _select_dtype(device: str) -> torch.dtype:
+        """Pick a stable dtype for the selected runtime device."""
         requested = os.environ.get("COFOS_DTYPE", "auto").strip().lower()
         aliases = {
             "fp32": torch.float32,
@@ -247,10 +265,12 @@ class StudentAdapterInference:
         if requested not in {"", "auto"}:
             logging.warning("Ignoring unsupported COFOS_DTYPE=%s; using auto", requested)
 
-        if torch.cuda.is_available():
+        if device == "cuda":
             if torch.cuda.is_bf16_supported():
                 return torch.bfloat16
             return torch.float16
+        if device == "mps":
+            return torch.float16
         return torch.float32
 
     def load_model(self) -> None:
@@ -277,15 +297,30 @@ class StudentAdapterInference:
         # Decoder-only models must be left-padded for correct generation.
         self.tokenizer.padding_side = "left"
 
-        dtype = self._select_dtype()
-        logging.info("Loading model weights with dtype=%s", dtype)
-        model = AutoModelForCausalLM.from_pretrained(
-            self.model_path_str,
-            device_map="auto",
-            torch_dtype=dtype,
-            low_cpu_mem_usage=True,
-            trust_remote_code=True,
-        )
+        device = self._select_device()
+        dtype = self._select_dtype(device)
+        logging.info("Loading model weights with dtype=%s on device=%s", dtype, device)
+        load_kwargs = {
+            "torch_dtype": dtype,
+            "low_cpu_mem_usage": True,
+            "trust_remote_code": True,
+        }
+        if device == "cuda":
+            load_kwargs["device_map"] = "auto"
+        try:
+            model = AutoModelForCausalLM.from_pretrained(self.model_path_str, **load_kwargs)
+            if device in {"mps", "cpu"}:
+                model = model.to(torch.device(device))
+        except ValueError as exc:
+            if "offload the whole model to the disk" in str(exc):
+                raise RuntimeError(
+                    "COFOS could not fit the 9B model in available memory. "
+                    "CUDA machines use automatic GPU placement; on Mac, COFOS avoids "
+                    "Transformers disk offload and tries direct MPS loading instead. "
+                    "Use a machine with more unified memory/GPU memory, or set "
+                    "COFOS_DEVICE=cpu for a slower CPU-only attempt."
+                ) from exc
+            raise
         logging.info("Using merged model without a separate LoRA adapter")
 
         model.generation_config.use_cache = True