From 7c5c0a41d71194bf2d9b30907a8b86c88a892942 Mon Sep 17 00:00:00 2001
From: expelledboy <102334+expelledboy@users.noreply.github.com>
Date: Sat, 23 Mar 2024 12:10:02 +0200
Subject: [PATCH] Add direnv with reproducible nix flake dev env

---
 .env.public |  1 +
 .envrc      |  4 ++++
 .gitignore  |  3 +++
 Justfile    | 14 ++++++++++++
 README.md   | 12 +++++++++++
 flake.lock  | 61 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 flake.nix   | 48 +++++++++++++++++++++++++++++++++++++++++
 7 files changed, 143 insertions(+)
 create mode 100644 .env.public
 create mode 100644 .envrc
 create mode 100644 Justfile
 create mode 100644 flake.lock
 create mode 100644 flake.nix

diff --git a/.env.public b/.env.public
new file mode 100644
index 0000000..e914361
--- /dev/null
+++ b/.env.public
@@ -0,0 +1 @@
+GROK_MAGNET_LINK=magnet:?xt=urn:btih:5f96d43576e3d386c9ba65b883210a393b68210e&tr=https%3A%2F%2Facademictorrents.com%2Fannounce.php&tr=udp%3A%2F%2Ftracker.coppersurfer.tk%3A6969&tr=udp%3A%2F%2Ftracker.opentrackr.org%3A1337%2Fannounce
diff --git a/.envrc b/.envrc
new file mode 100644
index 0000000..50ed70c
--- /dev/null
+++ b/.envrc
@@ -0,0 +1,4 @@
+dotenv .env.public
+dotenv_if_exists
+use flake
+layout python
diff --git a/.gitignore b/.gitignore
index 24d0d7e..5cfe4cb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,5 @@
+/.direnv/
+/.venv/
 checkpoints/*
 !checkpoints/README.md
+__pycache__/
diff --git a/Justfile b/Justfile
new file mode 100644
index 0000000..a56fe97
--- /dev/null
+++ b/Justfile
@@ -0,0 +1,14 @@
+[private]
+@default:
+	just --list --unsorted
+
+# Run the sample prompt against the model
+test:
+	python run.py
+
+# Download the weights for the Grok model
+download-weights:
+	transmission-cli \
+		--download-dir ./checkpoints \
+		$GROK_MAGNET_LINK
+	ln -s ./checkpoints/grok-1/ckpt-0 ./checkpoints/ckpt-0
diff --git a/README.md b/README.md
index f501a07..8a90bab 100644
--- a/README.md
+++ b/README.md
@@ -49,6 +49,18 @@ pip install huggingface_hub[hf_transfer]
 huggingface-cli download xai-org/grok-1 --repo-type model --include ckpt-0/* --local-dir checkpoints --local-dir-use-symlinks False
 ```
 
+# Reproducible Environment
+
+Those familiar with `nix` and `direnv` can use a reproducible development environment.
+
+```sh
+cd ~/src/grok-1
+direnv allow .
+just --list
+just download-weights
+just test
+```
+
 # License
 
 The code and associated Grok-1 weights in this release are licensed under the
diff --git a/flake.lock b/flake.lock
new file mode 100644
index 0000000..67d2ad1
--- /dev/null
+++ b/flake.lock
@@ -0,0 +1,61 @@
+{
+  "nodes": {
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1710146030,
+        "narHash": "sha256-SZ5L6eA7HJ/nmkzGG7/ISclqe6oZdOZTNoesiInkXPQ=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "b1d9ab70662946ef0850d488da1c9019f3a9752a",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1711124224,
+        "narHash": "sha256-l0zlN/3CiodvWDtfBOVxeTwYSRz93muVbXWSpaMjXxM=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "56528ee42526794d413d6f244648aaee4a7b56c0",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-23.11",
+        "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
+}
diff --git a/flake.nix b/flake.nix
new file mode 100644
index 0000000..0639b38
--- /dev/null
+++ b/flake.nix
@@ -0,0 +1,48 @@
+{
+  inputs = {
+    nixpkgs.url = "github:nixos/nixpkgs/nixos-23.11";
+    flake-utils.url = "github:numtide/flake-utils";
+  };
+
+  outputs = { self, nixpkgs, flake-utils }: flake-utils.lib.eachDefaultSystem (system:
+    let
+      pkgs = import nixpkgs { inherit system; };
+      python = pkgs.python311;
+    in
+    {
+      devShells.default = pkgs.mkShell {
+
+        # build-time
+        nativeBuildInputs = with pkgs; [
+          bashInteractive
+        ];
+
+        # run-time
+        buildInputs = (with pkgs; [
+          transmission
+          just
+        ]) ++ (with python.pkgs; [
+          setuptools
+          wheel
+          venvShellHook
+          pylint
+        ]);
+
+        # python setup
+        src = null;
+        venvDir = ".venv";
+        postVenv = ''
+          unset SOURCE_DATE_EPOCH
+        '';
+        postShellHook = ''
+          # python setup
+          unset SOURCE_DATE_EPOCH
+          unset LD_PRELOAD
+          PYTHONPATH=$PWD/$venvDir/${python.sitePackages}:$PYTHONPATH
+          pip install --require-virtualenv -r requirements.txt | grep -v 'already satisfied'
+        '';
+
+      };
+    }
+  );
+}