Add garden client and replace cluster implementation

lilac/auth.py

@@ -79,6 +79,9 @@ class AuthenticationInfo(BaseModel):
 
 def has_garden_credentials() -> bool:
   """Returns whether the user has Garden credentials."""
+  # TODO: more granular checks based on user permissions
+  if env('LILAC_API_KEY') is not None:
+    return True
   config = modal.config.Config().to_dict()
   return (
     'token_secret' in config

lilac/data/clustering.py

@@ -3,11 +3,11 @@
 import itertools
 from typing import Callable, Iterator, Optional, Union, cast
 
-import modal
 import numpy as np
 from tqdm import tqdm
 
-from ..batch_utils import compress_docs, flatten_path_iter
+from .. import garden_client
+from ..batch_utils import flatten_path_iter
 from ..dataset_format import DatasetFormatInputSelector
 from ..embeddings.jina import JinaV2Small
 from ..schema import (
@@ -311,11 +311,7 @@ def _hdbscan_cluster(
 ) -> Iterator[Item]:
   """Cluster docs with HDBSCAN."""
   if use_garden:
-    remote_fn = modal.Function.lookup('cluster', 'Cluster.cluster').remote
-    with DebugTimer('Compressing docs for clustering remotely'):
-      gzipped_docs = compress_docs(list(docs))
-    response = remote_fn({'gzipped_docs': gzipped_docs, 'min_cluster_size': min_cluster_size})
-    yield from response['clusters']
+    yield from garden_client.cluster(list(docs), min_cluster_size=min_cluster_size)
 
   if task_info:
     task_info.message = 'Computing embeddings'

lilac/env.py

@@ -56,6 +56,9 @@ class LilacEnvironment(BaseModel):
   COHERE_API_KEY: str = PydanticField(
     description='The Cohere API key, used for computing `cohere` embeddings.'
   )
+  LILAC_API_KEY: str = PydanticField(
+    description='The Lilac API key, used for running Lilac Garden computations.'
+  )
 
   # HuggingFace demo.
   HF_ACCESS_TOKEN: str = PydanticField(

lilac/garden_client.py

@@ -0,0 +1,47 @@
+"""Client code for sending requests to Lilac Garden."""
+import base64
+import functools
+import io
+import json
+from typing import Any, Callable, Iterator
+
+import numpy as np
+import requests
+
+from .env import env
+from .schema import Item
+from .utils import DebugTimer
+
+GARDEN_FRONT_GATE_URL = 'https://lilacai--front-gate-fastapi-app-dev.modal.run'
+GARDEN_ENCODING_SCHEME_HEADER = 'X-Lilac-EncodingScheme'
+
+
+def _decode_b64_npy(b: bytes) -> np.ndarray:
+  return np.load(io.BytesIO(base64.b64decode(b)))
+
+
+DECODERS: dict[str, Callable[[bytes], Item]] = {'b64-npy': _decode_b64_npy, 'json': json.loads}
+
+
+def _call_garden(endpoint_name: str, docs: list[Any], **kwargs: Any) -> Iterator[Item]:
+  lilac_api_key = env('LILAC_API_KEY')
+
+  with DebugTimer('Running garden endpoint %s' % endpoint_name):
+    with requests.post(
+      GARDEN_FRONT_GATE_URL + '/' + endpoint_name,
+      data=json.dumps(docs),
+      params={k: str(v) for k, v in kwargs.items()},
+      headers={
+        'Authorization': 'Bearer %s' % lilac_api_key,
+        'X-Lilac-RowCount': str(len(docs)),
+      },
+      stream=True,
+    ) as response:
+      if response.status_code > 299:
+        raise requests.HTTPError(response.text)
+      decoder = DECODERS[response.headers[GARDEN_ENCODING_SCHEME_HEADER]]
+      for line in response.iter_lines():
+        yield decoder(line)
+
+
+cluster = functools.partial(_call_garden, 'cluster')