Push

src/data/duckdb_utils.py

@@ -1,11 +1,15 @@
 """Utils for duckdb."""
 import duckdb
 
-from ..config import CONFIG
+from ..config import CONFIG, data_path
 
 
 def duckdb_gcs_setup(con: duckdb.DuckDBPyConnection) -> str:
   """Setup DuckDB for GCS."""
+  con.execute(f"""
+    SET extension_directory='{data_path()}';
+  """)
+
   con.install_extension('httpfs')
   con.load_extension('httpfs')
 

src/data/sources/json_source.py

@@ -6,7 +6,6 @@ import pandas as pd
 from pydantic import Field as PydanticField
 from typing_extensions import override
 
-from ...config import data_path
 from ...schema import Item
 from ...utils import download_http_files
 from ..duckdb_utils import duckdb_gcs_setup
@@ -39,7 +38,6 @@ class JSONDataset(Source):
     # DuckDB expects s3 protocol: https://duckdb.org/docs/guides/import/s3_import.html.
     s3_filepaths = [path.replace('gs://', 's3://') for path in filepaths]
 
-    con.execute(f"""SET extension_directory='{data_path()}';""")
     # NOTE: We use duckdb here to increase parallelism for multiple files.
     self._df = con.execute(f"""
       {duckdb_gcs_setup(con)}

src/embeddings/sbert.py

@@ -6,6 +6,7 @@ import torch
 from sentence_transformers import SentenceTransformer
 from typing_extensions import override
 
+from ..config import data_path
 from ..schema import Item, RichData
 from ..signals.signal import TextEmbeddingSignal
 from ..signals.splitters.chunk_splitter import split_text
@@ -32,7 +33,8 @@ def _sbert() -> tuple[Optional[str], SentenceTransformer]:
     preferred_device = 'mps'
   elif not torch.backends.mps.is_built():
     log('MPS not available because the current PyTorch install was not built with MPS enabled.')
-  return preferred_device, SentenceTransformer(MODEL_NAME, device=preferred_device)
+  return preferred_device, SentenceTransformer(
+    MODEL_NAME, device=preferred_device, cache_folder=data_path())
 
 
 def _optimal_batch_size(preferred_device: Optional[str]) -> int: