made dependancies pure

.envrc

@@ -1 +1 @@
-use nix
+use flake

.gitignore

@@ -1,2 +1,3 @@
 unsloth_compiled_cache
 outputs_qwen35
+__marimo__

README.md

@@ -2,11 +2,13 @@
 
 ## Overview
 
-### `shell.nix`
+### `flake.nix`
-Provides a development environment compatible within Nix. It handles:
+Provides a pure development environment via Nix Flakes. It handles:
-- Installing PyTorch optimized for AMD ROCm (`rocm7.2`).
+- Configures Python 3.13 with ROCm support enabled.
-- Installing `unsloth` and `unsloth-zoo` for efficient fine-tuning.
+- Provides PyTorch and specialized ML libraries (transformers, accelerate).
-- Installing `marimo` and `ipython` as QOL.
+- Packages `unsloth` and `unsloth-zoo` with necessary patches for compatibility.
+- Includes `marimo` and `ipython` for interactive development.
+- Automatically managed via `direnv` (`use flake`) or `nix develop`.
 
 ### `train.py`
 A `marimo` script that executes the fine-tuning process:

flake.lock

@@ -0,0 +1,61 @@
+{
+  "nodes": {
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1774386573,
+        "narHash": "sha256-4hAV26quOxdC6iyG7kYaZcM3VOskcPUrdCQd/nx8obc=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "46db2e09e1d3f113a13c0d7b81e2f221c63b8ce9",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-unstable",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,85 @@
+{
+  description = "Minimal example training a Qwen3.5 0.8B lora on AMD RYZEN AI MAX+ 395 APU";
+
+  inputs = {
+    nixpkgs.url = "github:nixos/nixpkgs?ref=nixos-unstable";
+    flake-utils.url = "github:numtide/flake-utils";
+  };
+
+  outputs = {
+    self,
+    nixpkgs,
+    flake-utils,
+  }:
+    flake-utils.lib.eachDefaultSystem (
+      system: let
+        pkgs = import nixpkgs {
+          inherit system;
+          config = {
+            allowUnfree = true;
+            rocmSupport = true;
+          };
+        };
+
+        python = pkgs.python313.override {
+          packageOverrides = self: super: {
+            datasets = super.datasets.overridePythonAttrs (oldAttrs: rec {
+              version = "4.3.0";
+              src = pkgs.fetchFromGitHub {
+                owner = "huggingface";
+                repo = "datasets";
+                tag = version;
+                hash = "sha256-3rDSHAMwoe9CkRLs3PDXSw2ONUrUWyBSZFpzk2C1A3A=";
+              };
+            });
+            unsloth-zoo = super.unsloth-zoo.overridePythonAttrs (oldAttrs: {
+              pythonImportsCheck = [];
+            });
+            unsloth = super.unsloth.overridePythonAttrs (oldAttrs: {
+              dependencies =
+                oldAttrs.dependencies
+                ++ [
+                  super.pydantic
+                  super.nest-asyncio
+                ];
+              pythonRelaxDeps =
+                oldAttrs.pythonRelaxDeps
+                ++ [
+                  "trl"
+                ];
+              postPatch = ''
+                # Relax setuptools version constraint in pyproject.toml
+                sed -i 's/setuptools==80\.9\.0/setuptools>=80.9/g' pyproject.toml || true
+
+                # Relax setuptools-scm version constraint in pyproject.toml
+                sed -i 's/setuptools-scm==9\.2\.0/setuptools-scm>=9.2/g' pyproject.toml || true
+              '';
+            });
+          };
+        };
+
+        pythonEnv = python.withPackages (ps:
+          with ps; [
+            ## Base ML/DS libs
+            torch
+            torchvision
+            torchaudio
+            transformers
+            accelerate
+
+            ## Unsloth
+            unsloth
+
+            ## QOL
+            marimo
+            ipython
+          ]);
+      in {
+        devShell = pkgs.mkShell {
+          packages = [
+            pythonEnv
+          ];
+        };
+      }
+    );
+}

shell.nix

@@ -1,21 +0,0 @@
-{nixpkgs ? import <nixpkgs> {}}:
-nixpkgs.mkShell {
-  nativeBuildInputs = with nixpkgs; [
-    ruff
-    python3
-  ];
-
-  LD_LIBRARY_PATH = "${nixpkgs.stdenv.cc.cc.lib}/lib:${nixpkgs.zstd.out}/lib:${nixpkgs.zlib.out}/lib";
-  shellHook = ''
-    if [[ ! -d ".venv" ]]; then
-      python -m venv .venv
-      source .venv/bin/activate
-      pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2/ --upgrade --force-reinstall
-      pip install "unsloth[amd] @ git+https://github.com/unslothai/unsloth"
-      pip install "unsloth-zoo[main] @ git+https://github.com/unslothai/unsloth-zoo"
-      pip install marimo ipython
-    else
-      source .venv/bin/activate
-    fi
-  '';
-}