dotfiles/nix/home-manager/progs/gpt4all/gpt4all.nix

{ pkgs, lib, ... }:
let
  models = [
    {
      name = "Qwen2.5-14B-Instruct-Q4_K_S.gguf";
      context_length = "32768";
      gen_length = "8192";
      source = pkgs.fetchurl {
        url = "https://huggingface.co/bartowski/Qwen2.5-14B-Instruct-GGUF/resolve/main/Qwen2.5-14B-Instruct-Q4_K_S.gguf?download=true";
        sha256 = "E1CmWUhMMbTXEjIRczzA3rSrVuR8qOL8BLagw7LiyZk=";
      };
    }
    {
      name = "Qwen2.5-7B-Instruct-Q6_K_L.gguf";
      context_length = "32768";
      gen_length = "8192";
      source = pkgs.fetchurl {
        url = "https://huggingface.co/bartowski/Qwen2.5-7B-Instruct-GGUF/resolve/main/Qwen2.5-7B-Instruct-Q6_K_L.gguf?download=true";
        sha256 = "thEXN06T/UVGfzdB83jlgpG7kuTzZtz1ZUAdupAnErM=";
      };
    }
  ];

  # stolen from: https://stackoverflow.com/a/42398526
  optimizeWithFlags =
    pkg: flags:
    pkgs.lib.overrideDerivation pkg (
      old:
      let
        newflags = pkgs.lib.foldl' (acc: x: "${acc} ${x}") "" flags;
        oldflags = if (pkgs.lib.hasAttr "NIX_CFLAGS_COMPILE" old) then "${old.NIX_CFLAGS_COMPILE}" else "";
      in
      {
        NIX_CFLAGS_COMPILE = "${oldflags} ${newflags}";
        stdenv = pkgs.clangStdenv;
      }
    );

  model_files = builtins.listToAttrs (
    map (f: {
      name = ".local/share/nomic.ai/GPT4All/${f.name}";
      value = {
        source = f.source;
      };
    }) models
  );

  gpt4all_package = (
    optimizeWithFlags
      (pkgs.gpt4all.overrideAttrs (old: {
        # https://github.com/NixOS/nixpkgs/pull/344001 3.2.1 -> 3.3.0
        version = "3.3.0";
        src = pkgs.fetchFromGitHub {
          fetchSubmodules = true;
          owner = "nomic-ai";
          repo = "gpt4all";
          rev = "HEAD";
          sha256 = "VTsdoXig/BEgPMkaEkHpEHwQ835aU2pUPhIxfi8SIhQ=";
        };

        patches = [
          ./gpt4all-HEAD-embeddings-model.patch
          ./gpt4all-HEAD-disable-settings-err.patch
        ];
      }))
      # compile flags
      [
        "-Ofast"
        "-march=native"
        "-mtune=native"
        "-fno-protect-parens"
        "-fno-finite-math-only" # https://github.com/ggerganov/llama.cpp/pull/7154#issuecomment-2143844461
      ]
  );

in
{
  home.packages = [
    gpt4all_package
  ];

  home.file = lib.recursiveUpdate {
    ".config/nomic.ai/GPT4All.ini".text =
      let
        system_prompt = "You are an expert AI assistant that explains your reasoning step by step. For each step, provide a title that describes what you're doing in that step, along with the content. Decide if you need another step or if you're ready to give the final answer. USE AS MANY REASONING STEPS AS POSSIBLE. AT LEAST 3. BE AWARE OF YOUR LIMITATIONS AS AN LLM AND WHAT YOU CAN AND CANNOT DO. EXPLORE ALTERNATE ANSWERS AND CONSIDER THAT YOUR ANSWER MAY BE WRONG. IDENTIFY POSSIBLE ERRORS IN YOUR REASONING AND WHERE SUCH ERRORS MAY BE. FULLY TEST ALL OTHER POSSIBILITIES. YOU CAN BE WRONG. WHEN YOU SAY YOU ARE RE-EXAMINING, ACTUALLY RE-EXAMINE, AND USE ANOTHER APPROACH TO DO SO. DO NOT JUST SAY YOU ARE RE-EXAMINING. SHOW ALL YOUR WORK. USE AT LEAST 3 METHODS TO DERIVE THE ANSWER. USE BEST PRACTICES. WORK FROM FIRST PRINCIPLES TO CREATE YOUR ANSWER.

YOUR REASONING IS NOT SHOWN TO THE USER BUT IS SUBJECT TO HIGH LEVELS OF SCRUTINY. WHEN YOU DETERMINE THAT YOU ARE READY TO GIVE A FINAL ANSWER TO THE USER GIVEN YOUR REASONING AND STEP-BY-STEP WORK, CREATE A SECTION WITH A \"Final Answer\" TITLE, THIS WILL BE THE ONLY RESPONSE SHOWN TO THE END USER, BE SURE TO INCLUDE ALL IMPORTANT INFORMATION AND DETAILS WHICH THE USER WILL FIND USEFUL TO AID THEIR QUESTION.";
      in
      ''
        [General]
        chatTheme=Dark
        height=940
        suggestionMode=Off
        threadCount=8
        uniqueId=7096f2d2-448d-4272-a132-d37e77f8a781
        userDefaultModel=${(builtins.elemAt models 0).name}
        width=1472
        x=0
        y=0

        [download]
        lastVersionStarted=${gpt4all_package.version}
      ''
      + (lib.concatStrings (
        map (model: ''
          [model-${model.name}]
          contextLength=${model.context_length}
          filename=${model.name}
          maxLength=${model.gen_length}
          promptBatchSize=256
          promptTemplate=<|im_start|>user\n%1<|im_end|>\n<|im_start|>assistant\n
          systemPrompt="<|im_start|>system\n${
            # replace newlines with the string "\n" for gpt4all to properly parse
            builtins.replaceStrings [ "\n" ] [ "\\n" ] system_prompt
          }<|im_end|>
          \n"
        '') models
      ))
      + ''
        [network]
        isActive=true
        usageStatsActive=true
      '';
  } model_files;
}