Merge pull request #2 from tillahoffmann/action

Add GitHub action.

Author: tillahoffmann

.github/workflows/main.yaml

@@ -0,0 +1,29 @@
+name: trecs
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["main"]
+
+jobs:
+  build:
+    name: Build
+    runs-on: "ubuntu-latest"
+    steps:
+      - uses: "actions/checkout@v4"
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+          cache-dependency-glob: uv.lock
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version-file: .python-version
+      - name: Install project
+        run: uv sync --all-groups
+      - name: Lint the code
+        run: uv run black --check .
+      - name: Run the tests
+        run: uv run pytest -v tests

README.md

@@ -1,4 +1,4 @@
-# 🎶 trecs: Transformer-Based Playlist Recommendation
+# 🎶 trecs: Transformer-Based Playlist Recommendation [![trecs](https://github.com/tillahoffmann/trecs/actions/workflows/main.yaml/badge.svg)](https://github.com/tillahoffmann/trecs/actions/workflows/main.yaml)
 
 This repository implements a transformer architecture to complete Spotify playlists given a set of seed tracks. The model is trained on the [Million Playlist Dataset](https://www.aicrowd.com/challenges/spotify-million-playlist-dataset-challenge) (MPD).
 

src/trecs/data/util.py

@@ -90,6 +90,19 @@ def _get_conn(self) -> sqlite3.Connection:
             self._local.conn = sqlite3.connect(self.conn)
             return self._local.conn
 
+    def _close_conn(self) -> None:
+        # Close the connection if there is one.
+        if isinstance(self.conn, sqlite3.Connection):
+            self.conn.close()
+        elif isinstance(self._local.conn, sqlite3.Connection):
+            self._local.conn.close()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, *_) -> None:
+        self._close_conn()
+
     @property
     def idx(self) -> list[Any]:
         if self._idx is None:

src/trecs/scripts/predict.py

@@ -1,4 +1 @@
 import argparse
-
-
-

src/trecs/scripts/train.py

@@ -49,7 +49,9 @@ def restore(
     rngs: nnx.Rngs | None = None,
 ) -> int:
     flax_nodes = {"model": model, "optimizer": optimizer, "rngs": rngs}
-    flax_states = {key: nnx.state(value) for key, value in flax_nodes.items() if value is not None}
+    flax_states = {
+        key: nnx.state(value) for key, value in flax_nodes.items() if value is not None
+    }
     data_iterators = data_iterators or {}
     restored = checkpoint_manager.restore(
         step,