python3Packages.ai-edge-litert: init at 2.1.0

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 ];
+  };
+}

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="
+      }
+    }
+  }
+}

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"^(?P<target>macosx|manylinux)[\d+_]+(?P<arch>x86_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<major>\d)(?P<minor>\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<pname>\w+)-{re.escape(version)}-(?P<python>\w+)-(?P<dist>\w+)-(?P<platform>\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")

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 { };