1{
2 lib,
3 buildPythonPackage,
4 einops,
5 fetchFromGitHub,
6 hyper-connections,
7 pytestCheckHook,
8 setuptools,
9 torch,
10}:
11
12buildPythonPackage rec {
13 pname = "local-attention";
14 version = "1.11.2";
15 pyproject = true;
16
17 src = fetchFromGitHub {
18 owner = "lucidrains";
19 repo = "local-attention";
20 tag = version;
21 hash = "sha256-2gBPALJAflLf7Y8L5wnNw4fHcvIOKjOncLsebkhrYkU=";
22 };
23
24 postPatch = ''
25 substituteInPlace setup.py \
26 --replace-fail "'pytest-runner'," ""
27 '';
28
29 build-system = [ setuptools ];
30
31 dependencies = [
32 einops
33 hyper-connections
34 torch
35 ];
36
37 nativeCheckInputs = [ pytestCheckHook ];
38
39 pythonImportsCheck = [ "local_attention" ];
40
41 meta = {
42 description = "Module for local windowed attention for language modeling";
43 homepage = "https://github.com/lucidrains/local-attention";
44 changelog = "https://github.com/lucidrains/local-attention/releases/tag/${src.tag}";
45 license = lib.licenses.mit;
46 maintainers = with lib.maintainers; [ fab ];
47 };
48}