87 lines
2.0 KiB
Nix
87 lines
2.0 KiB
Nix
{
|
|
inputs = {
|
|
tgi-nix.url = "github:danieldk/tgi-nix";
|
|
nixpkgs.follows = "tgi-nix/nixpkgs";
|
|
flake-utils.url = "github:numtide/flake-utils";
|
|
rust-overlay = {
|
|
url = "github:oxalica/rust-overlay";
|
|
inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
|
|
};
|
|
};
|
|
outputs =
|
|
{
|
|
self,
|
|
nixpkgs,
|
|
flake-utils,
|
|
rust-overlay,
|
|
tgi-nix,
|
|
}:
|
|
flake-utils.lib.eachDefaultSystem (
|
|
system:
|
|
let
|
|
config = {
|
|
allowUnfree = true;
|
|
cudaSupport = true;
|
|
};
|
|
pkgs = import nixpkgs {
|
|
inherit config system;
|
|
overlays = [
|
|
rust-overlay.overlays.default
|
|
tgi-nix.overlay
|
|
];
|
|
};
|
|
in
|
|
{
|
|
devShells.default =
|
|
with pkgs;
|
|
mkShell {
|
|
buildInputs =
|
|
[
|
|
openssl.dev
|
|
pkg-config
|
|
(rust-bin.stable.latest.default.override {
|
|
extensions = [
|
|
"rust-analyzer"
|
|
"rust-src"
|
|
];
|
|
})
|
|
]
|
|
++ (with python3.pkgs; [
|
|
venvShellHook
|
|
pip
|
|
|
|
einops
|
|
fbgemm-gpu
|
|
flash-attn
|
|
flash-attn-layer-norm
|
|
flash-attn-rotary
|
|
grpc-interceptor
|
|
grpcio-reflection
|
|
grpcio-status
|
|
hf-transfer
|
|
loguru
|
|
marlin-kernels
|
|
opentelemetry-api
|
|
opentelemetry-exporter-otlp
|
|
opentelemetry-instrumentation-grpc
|
|
opentelemetry-semantic-conventions
|
|
peft
|
|
tokenizers
|
|
torch
|
|
transformers
|
|
vllm
|
|
]);
|
|
|
|
venvDir = "./.venv";
|
|
|
|
postVenv = ''
|
|
unset SOURCE_DATE_EPOCH
|
|
'';
|
|
postShellHook = ''
|
|
unset SOURCE_DATE_EPOCH
|
|
'';
|
|
};
|
|
}
|
|
);
|
|
}
|