Upload app.py with huggingface_hub

app.py

@@ -20,27 +20,52 @@ HUB_PATH = os.environ.get(
     "training-runs/sft-chat-v1-context-mix-8b-20260803-080404",
 )
 SCRIPT_REPO = os.environ.get("SCRIPT_REPO", "ProjectScugnizz/scugnizz-llama-training")
+# Space volumes: checkpoint subfolder → /models, trainer → /code
+MODEL_MOUNT = Path(os.environ.get("MODEL_MOUNT", "/models"))
+CODE_MOUNT = Path(os.environ.get("CODE_MOUNT", "/code"))
 SYSTEM = os.environ.get(
     "SYSTEM_PROMPT",
     "You are a helpful assistant. Answer clearly and correctly.",
 )
 
-_CODE = Path("/tmp/scugnizz-code")
-_MODEL = Path("/tmp/scugnizz-model")
+_CODE_FALLBACK = Path("/tmp/scugnizz-code")
+_MODEL_FALLBACK = Path("/tmp/scugnizz-model")
 _state = {"model": None, "tok": None, "dev": None, "error": None}
 
 
-def _load_train_module():
+def _resolve_code_dir() -> Path:
+    if (CODE_MOUNT / "scugnizz-llama.py").is_file():
+        return CODE_MOUNT
     snapshot_download(
         SCRIPT_REPO,
-        local_dir=str(_CODE),
+        local_dir=str(_CODE_FALLBACK),
         allow_patterns=["scugnizz-llama.py", "sft_data.py"],
     )
-    path = _CODE / "scugnizz-llama.py"
+    return _CODE_FALLBACK
+
+
+def _resolve_model_dir() -> Path:
+    # Volume mounts either the checkpoint root (/models/model_final.pt) or full repo
+    if (MODEL_MOUNT / "model_final.pt").is_file():
+        return MODEL_MOUNT
+    nested = MODEL_MOUNT / HUB_PATH.strip("/")
+    if (nested / "model_final.pt").is_file():
+        return nested
+    path = HUB_PATH.strip("/")
+    snapshot_download(
+        HUB_REPO,
+        local_dir=str(_MODEL_FALLBACK),
+        allow_patterns=[f"{path}/*", f"{path}/tokenizer/*"],
+    )
+    return _MODEL_FALLBACK / path
+
+
+def _load_train_module():
+    code = _resolve_code_dir()
+    path = code / "scugnizz-llama.py"
     spec = importlib.util.spec_from_file_location("scugnizz_train", path)
     mod = importlib.util.module_from_spec(spec)
-    # sft_data lives beside the trainer file
-    sys.path.insert(0, str(_CODE))
+    sys.path.insert(0, str(code))
     sys.modules["scugnizz_train"] = mod
     spec.loader.exec_module(mod)
     return mod
@@ -101,13 +126,7 @@ def load_model():
         raise RuntimeError(_state["error"])
     try:
         mod = _load_train_module()
-        path = HUB_PATH.strip("/")
-        snapshot_download(
-            HUB_REPO,
-            local_dir=str(_MODEL),
-            allow_patterns=[f"{path}/*", f"{path}/tokenizer/*"],
-        )
-        model_dir = _MODEL / path
+        model_dir = _resolve_model_dir()
         run_args = json.loads((model_dir / "args.json").read_text(encoding="utf-8"))
         dev = "cuda" if torch.cuda.is_available() else "cpu"
         tok = AutoTokenizer.from_pretrained(str(model_dir / "tokenizer"))