Switch splitter to new algorithm, fixing last blocking bug. (#796)

lilac/embeddings/cohere.py

@@ -7,7 +7,7 @@
 from ..env import env
 from ..schema import Item, RichData
 from ..signal import TextEmbeddingSignal
-from ..splitters.chunk_splitter import split_text
+from ..splitters.spacy_splitter import clustering_spacy_chunker
 from .embedding import compute_split_embeddings
 
 if TYPE_CHECKING:
@@ -54,6 +54,6 @@ def embed_fn(texts: list[str]) -> list[np.ndarray]:
       return self._model.embed(texts, truncate='END').embeddings
 
     docs = cast(Iterable[str], docs)
-    split_fn = split_text if self._split else None
+    split_fn = clustering_spacy_chunker if self._split else None
     yield from compute_split_embeddings(
       docs, COHERE_BATCH_SIZE, embed_fn, split_fn, num_parallel_requests=NUM_PARALLEL_REQUESTS)

lilac/embeddings/gte.py

@@ -5,7 +5,7 @@
 
 from ..schema import Item, RichData
 from ..signal import TextEmbeddingSignal
-from ..splitters.chunk_splitter import split_text
+from ..splitters.spacy_splitter import clustering_spacy_chunker
 from .embedding import compute_split_embeddings
 from .transformer_utils import get_model
 
@@ -48,7 +48,7 @@ def compute(self, docs: Iterable[RichData]) -> Iterable[Item]:
     """Call the embedding function."""
     batch_size, model = get_model(self._model_name, _OPTIMAL_BATCH_SIZES[self._model_name])
     embed_fn = model.encode
-    split_fn = split_text if self._split else None
+    split_fn = clustering_spacy_chunker if self._split else None
     docs = cast(Iterable[str], docs)
     yield from compute_split_embeddings(docs, batch_size, embed_fn=embed_fn, split_fn=split_fn)
 

lilac/embeddings/openai.py

@@ -8,7 +8,7 @@
 from ..env import env
 from ..schema import Item, RichData
 from ..signal import TextEmbeddingSignal
-from ..splitters.chunk_splitter import split_text
+from ..splitters.spacy_splitter import clustering_spacy_chunker
 from .embedding import compute_split_embeddings
 
 if TYPE_CHECKING:
@@ -63,6 +63,6 @@ def embed_fn(texts: list[str]) -> list[np.ndarray]:
       return [np.array(embedding['embedding'], dtype=np.float32) for embedding in response['data']]
 
     docs = cast(Iterable[str], docs)
-    split_fn = split_text if self._split else None
+    split_fn = clustering_spacy_chunker if self._split else None
     yield from compute_split_embeddings(
       docs, OPENAI_BATCH_SIZE, embed_fn, split_fn, num_parallel_requests=NUM_PARALLEL_REQUESTS)

lilac/embeddings/palm.py

@@ -8,7 +8,7 @@
 from ..env import env
 from ..schema import Item, RichData
 from ..signal import TextEmbeddingSignal
-from ..splitters.chunk_splitter import split_text
+from ..splitters.spacy_splitter import clustering_spacy_chunker
 from .embedding import compute_split_embeddings
 
 if TYPE_CHECKING:
@@ -57,6 +57,6 @@ def embed_fn(texts: list[str]) -> list[np.ndarray]:
       return [np.array(response['embedding'], dtype=np.float32)]
 
     docs = cast(Iterable[str], docs)
-    split_fn = split_text if self._split else None
+    split_fn = clustering_spacy_chunker if self._split else None
     yield from compute_split_embeddings(
       docs, PALM_BATCH_SIZE, embed_fn, split_fn, num_parallel_requests=NUM_PARALLEL_REQUESTS)

lilac/embeddings/sbert.py

@@ -5,7 +5,7 @@
 
 from ..schema import Item, RichData
 from ..signal import TextEmbeddingSignal
-from ..splitters.chunk_splitter import split_text
+from ..splitters.spacy_splitter import clustering_spacy_chunker
 from .embedding import compute_split_embeddings
 from .transformer_utils import get_model
 
@@ -33,6 +33,6 @@ def compute(self, docs: Iterable[RichData]) -> Iterable[Item]:
     """Call the embedding function."""
     batch_size, model = get_model(MINI_LM_MODEL, _OPTIMAL_BATCH_SIZES[MINI_LM_MODEL])
     embed_fn = model.encode
-    split_fn = split_text if self._split else None
+    split_fn = clustering_spacy_chunker if self._split else None
     docs = cast(Iterable[str], docs)
     yield from compute_split_embeddings(docs, batch_size, embed_fn=embed_fn, split_fn=split_fn)

lilac/schema.py

@@ -366,7 +366,8 @@ def lilac_span(start: int, end: int, metadata: dict[str, Any] = {}) -> Item:
 
 def lilac_embedding(start: int, end: int, embedding: Optional[np.ndarray]) -> Item:
   """Creates a lilac embedding item, representing a vector with a pointer to a slice of text."""
-  return lilac_span(start, end, {EMBEDDING_KEY: embedding})
+  # Cast to int; we've had issues where start/end were np.int64, which caused downstream sadness.
+  return lilac_span(int(start), int(end), {EMBEDDING_KEY: embedding})
 
 
 def _parse_field_like(field_like: object, dtype: Optional[Union[DataType, str]] = None) -> Field: