diff --git a/pkgs/development/python-modules/ai-edge-litert/default.nix b/pkgs/development/python-modules/ai-edge-litert/default.nix new file mode 100644 index 000000000000..71488fe46c7a --- /dev/null +++ b/pkgs/development/python-modules/ai-edge-litert/default.nix @@ -0,0 +1,87 @@ +{ + autoPatchelfHook, + buildPythonPackage, + fetchurl, + lib, + python, + pythonAtLeast, + stdenv, + + # dependencies + backports-strenum, + flatbuffers, + numpy, + protobuf, + tqdm, + typing-extensions, + + # optional-dependencies + lark, + ml-dtypes +}: + +let + release = builtins.fromJSON (builtins.readFile ./release.json); + platforms = release.src; + platform = + platforms.${stdenv.hostPlatform.system} + or (throw "ai-edge-litert: unsupported platform (${stdenv.hostPlatform.system})"); + pythonMajorMinor = lib.versions.majorMinor python.version; + source = + platform.${pythonMajorMinor} + or (throw "ai-edge-litert: unsupported python version (${pythonMajorMinor})"); +in + +buildPythonPackage { + pname = "ai-edge-litert"; + version = release.version; + format = "wheel"; + + src = fetchurl { + inherit (source) + url + hash + ; + }; + + nativeBuildInputs = [ autoPatchelfHook ]; + + dependencies = [ + backports-strenum + flatbuffers + numpy + protobuf + tqdm + typing-extensions + ]; + + optional-dependencies = { + model-utils = [ + lark + ml-dtypes + # TODO :xdsl + ]; + # TODO: npu-sdk + }; + + pythonRemoveDeps = lib.optionals (pythonAtLeast "3.12") [ + # https://github.com/google-ai-edge/LiteRT/pull/5298 + "backports.strenum" + ]; + + pythonImportsCheck = [ + "ai_edge_litert" + "ai_edge_litert.interpreter" + ]; + + passthru.updateScript = ./update.py; + + meta = { + description = "LiteRT is for mobile and embedded devices"; + downloadPage = "https://github.com/google-ai-edge/LiteRT"; + homepage = "https://www.tensorflow.org/lite/"; + license = lib.licenses.asl20; + platforms = lib.attrNames platforms; + maintainers = with lib.maintainers; [ hexa ]; + }; +} diff --git a/pkgs/development/python-modules/ai-edge-litert/release.json b/pkgs/development/python-modules/ai-edge-litert/release.json new file mode 100644 index 000000000000..d21b94acfeb9 --- /dev/null +++ b/pkgs/development/python-modules/ai-edge-litert/release.json @@ -0,0 +1,49 @@ +{ + "version": "2.1.0", + "src": { + "aarch64-linux": { + "3.10": { + "url": "https://files.pythonhosted.org/packages/a5/7a/6ea77fe405d76dcd1a95fceccd1018293603f551bf3645a382ef1fd470eb/ai_edge_litert-2.1.0-cp310-cp310-manylinux_2_27_aarch64.whl", + "hash": "sha256-eAZzWROFxGRqqCs42qLu2IvCY0bhqvvp0Nt/BJPWq1Q=" + }, + "3.11": { + "url": "https://files.pythonhosted.org/packages/82/e7/73fff29c7a95a4f35953a027bdb174dd9ac08adbe8c2c322dee26d8f1d3f/ai_edge_litert-2.1.0-cp311-cp311-manylinux_2_27_aarch64.whl", + "hash": "sha256-krJNJQf8z1aAwm5xLBViZtz2Sokt85Id1C9KfSZSdGI=" + }, + "3.12": { + "url": "https://files.pythonhosted.org/packages/61/f6/1aeccf3034f95ec680a2e25f33582d4adf02c30b3462357511d529a2cad2/ai_edge_litert-2.1.0-cp312-cp312-manylinux_2_27_aarch64.whl", + "hash": "sha256-4Gt9sF08WCT9S1biMY9oWN73GHgM2QMYYmtKqY23WDw=" + }, + "3.13": { + "url": "https://files.pythonhosted.org/packages/9e/b2/dde0f8043eb26f237c5958e44334d62fa90496fa51438f9c5179dec72d4c/ai_edge_litert-2.1.0-cp313-cp313-manylinux_2_27_aarch64.whl", + "hash": "sha256-eoB6MjbIz1ZeIxtuwzG+zjo/WhfDwl+I8+N8/oZo+NA=" + }, + "3.9": { + "url": "https://files.pythonhosted.org/packages/30/8b/531b230a5377b2253eff8235f10e72489bc207f7e5ed6e531670495e4fd4/ai_edge_litert-2.1.0-cp39-cp39-manylinux_2_27_aarch64.whl", + "hash": "sha256-qe3jwpMVGnSt+hK6yLjOh+wB3Fhuw86hIqxxRlDIQb8=" + } + }, + "x86_64-linux": { + "3.10": { + "url": "https://files.pythonhosted.org/packages/20/4a/5f7848c8d21df432f44ebd36097406f66d98762295a4f11239f1d913f347/ai_edge_litert-2.1.0-cp310-cp310-manylinux_2_27_x86_64.whl", + "hash": "sha256-0FSis6ZeLCtFcahO0zul2ojp0Xhy/YdTwXc/+3g1HLE=" + }, + "3.11": { + "url": "https://files.pythonhosted.org/packages/0d/d1/032cd1300a9b40da97e9f93781b3c2c5f104c22157695a3bb5bd5531821f/ai_edge_litert-2.1.0-cp311-cp311-manylinux_2_27_x86_64.whl", + "hash": "sha256-3PfXk3SrRnzsGfLtrdO+skOtF4M/jFrVm/dnEr8USbE=" + }, + "3.12": { + "url": "https://files.pythonhosted.org/packages/bb/16/0aface4ec3554304487d6436938684023c054828b4334f114f7db303b27a/ai_edge_litert-2.1.0-cp312-cp312-manylinux_2_27_x86_64.whl", + "hash": "sha256-SNfm6BuqNYHmtIkgpzq4lzAcxIqluJq9012fBt9jLUo=" + }, + "3.13": { + "url": "https://files.pythonhosted.org/packages/96/c7/53b1f2627dbb0b307c1637e5f2cdbdc5a5f4ffecbebe75c8977d46590b6a/ai_edge_litert-2.1.0-cp313-cp313-manylinux_2_27_x86_64.whl", + "hash": "sha256-QeJJDL8dJmZj1FEllzEXC4Gqpwt9/vbXGIaX5lxxLJU=" + }, + "3.9": { + "url": "https://files.pythonhosted.org/packages/85/b7/efd39c318ae7f73bd3de247a93501f1d748e60979ecd97c794e7a4a1d671/ai_edge_litert-2.1.0-cp39-cp39-manylinux_2_27_x86_64.whl", + "hash": "sha256-QicA/OOs8HK8DVjXhj7BZjHwix2t3vyHwAjebkzaMWc=" + } + } + } +} \ No newline at end of file diff --git a/pkgs/development/python-modules/ai-edge-litert/update.py b/pkgs/development/python-modules/ai-edge-litert/update.py new file mode 100755 index 000000000000..e792c3a06443 --- /dev/null +++ b/pkgs/development/python-modules/ai-edge-litert/update.py @@ -0,0 +1,99 @@ +#!/usr/bin/env nix-shell +#!nix-shell -i python -p "python3.withPackages (ps: [ ps.httpx ] )" + +import json +import httpx +import re +import subprocess +from pathlib import Path +from typing import Dict +from collections import defaultdict + + +def get_metadata(package) -> Dict: + response = httpx.get( + f"https://pypi.org/pypi/{package}/json", + headers={"user-agent": "nixpkgs/pypi-updater/unstable"}, + ) + return response.json() + + +def get_hash(url) -> str: + result = subprocess.run(["nix-prefetch-url", url], stdout=subprocess.PIPE) + base32_hash = result.stdout.decode().strip() + result = subprocess.run( + ["nix", "hash", "to-sri", "--type", "sha256", base32_hash], + stdout=subprocess.PIPE, + ) + sri_hash = result.stdout.decode().strip() + return sri_hash + + +def get_platform(platform: str) -> str: + result = re.match(r"^(?Pmacosx|manylinux)[\d+_]+(?Px86_64|aarch64)", platform) + if not result: + raise RuntimeError(f"Unable to parse platform string: {platform}") + + target = { + "macosx": "darwin", + "manylinux": "linux", + }[result.group("target")] + + arch = result.group("arch") + + return f"{arch}-{target}" + +def get_python_version(python: str) -> str: + result = re.match(r"^cp(?P\d)(?P\d+)$", python) + if not result: + raise RuntimeError(f"Unable to disect python compat tag: {python}") + + return f"{result.group('major')}.{result.group('minor')}" + + +def main(package: str): + metadata = get_metadata(package) + + info = metadata.get("info") + if info is None: + raise RuntimeError("Package metadata has no info attribute") + + version = info.get("version") + assert not info.get("yanked"), ( + f"Latest release was yanked: {info.get('yanked_reason')}" + ) + releases = metadata["releases"][version] + + out = { + "version": version, + "src": defaultdict(dict), + } + for release in releases: + if release.get("packagetype") != "bdist_wheel": + # the package expects the binary wheels + continue + + print(json.dumps(release, indent=2)) + + filename = release.get("filename") + result = re.match( + rf"(?P\w+)-{re.escape(version)}-(?P\w+)-(?P\w+)-(?P\w+)\.whl$", + filename, + ) + if not result: + raise RuntimeError(f"Unable to disect wheel filename: {filename}") + + platform = get_platform(result.group("platform")) + python_version = get_python_version(release["python_version"]) + + out["src"][platform][python_version] = { + "url": release.get("url"), + "hash": get_hash(release.get("url")), + } + + with open(Path(__file__).with_name("release.json"), "w") as fd: + json.dump(out, fd, indent=2) + + +if __name__ == "__main__": + main("ai-edge-litert") diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index caf81aa8d018..47d8186b449b 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -156,6 +156,8 @@ self: super: with self; { ahocorasick-rs = callPackage ../development/python-modules/ahocorasick-rs { }; + ai-edge-litert = callPackage ../development/python-modules/ai-edge-litert { }; + aigpy = callPackage ../development/python-modules/aigpy { }; ailment = callPackage ../development/python-modules/ailment { };