Renamed to LateSentenceTransformer

Author: x-tabdeveloping

turftopic/encoders/contextual.py turftopic/encoders/late_interaction.pyturftopic/encoders/contextual.py renamed to turftopic/encoders/late_interaction.py

@@ -18,7 +18,7 @@ def is_contextual(encoder):
 Lengths = list[int]
 
 
-class LateTransformer(SentenceTransformer):
+class LateSentenceTransformer(SentenceTransformer):
     def encode(
         self, sentences: Union[str, list[str], np.ndarray], *args, **kwargs
     ):

turftopic/late.py

@@ -3,6 +3,8 @@
 import numpy as np
 from sklearn.base import TransformerMixin
 
+from turftopic.encoders.contextual import Offsets
+
 Lengths = list[int]
 
 
@@ -62,7 +64,17 @@ def pool_flat(flat_repr: np.ndarray, lengths: Lengths, agg=np.mean):
     return np.stack(pooled)
 
 
-class TokenLevel(TransformerMixin):
+def get_document_chunks(
+    raw_documents: list[str], offsets: list[Offsets]
+) -> list[str]:
+    chunks = []
+    for doc, _offs in zip(raw_documents, offsets):
+        for start_char, end_char in _offs:
+            chunks.append(raw_documents[start_char, end_char])
+    return chunks
+
+
+class LateModel(TransformerMixin):
     def __init__(
         self,
         model: TransformerMixin,
@@ -74,63 +86,18 @@ def __init__(
         self.pooling = pooling
 
     def transform(
-        self, raw_documents: list[str], embeddings: list[np.ndarray] = None
-    ):
-        if embeddings is None:
-            embeddings = self.model.encoder.encode_tokens(
-                raw_documents, batch_size=self.batch_size
-            )
-        flat_embeddings, lengths = flatten_repr(embeddings)
-        out_array = self.model.transform(
-            raw_documents, embeddings=flat_embeddings
-        )
-        if self.pooling is None:
-            return unflatten_repr(out_array, lengths)
-        else:
-            return pool_flat(out_array, lengths)
-
-    def fit_transform(
         self,
         raw_documents: list[str],
-        y=None,
         embeddings: list[np.ndarray] = None,
+        offsets: list[Offsets] = None,
     ):
-        if embeddings is None:
-            embeddings = self.model.encoder.encode_tokens(
+        if (embeddings is None) or (offsets is None):
+            embeddings, offsets = self.model.encoder.encode_tokens(
                 raw_documents, batch_size=self.batch_size
             )
         flat_embeddings, lengths = flatten_repr(embeddings)
-        out_array = self.model.fit_transform(
-            raw_documents, y, embeddings=flat_embeddings
-        )
-        if self.pooling is None:
-            return unflatten_repr(out_array, lengths)
-        else:
-            return pool_flat(out_array, lengths)
-
-
-class Windowed(TransformerMixin):
-    def __init__(
-        self,
-        model: TransformerMixin,
-        batch_size: int = 32,
-        pooling: Optional[Callable] = None,
-    ):
-        self.model = model
-        self.batch_size = batch_size
-        self.pooling = pooling
-
-    def transform(
-        self, raw_documents: list[str], embeddings: list[np.ndarray] = None
-    ):
-        if embeddings is None:
-            embeddings = self.model.encoder.encode_tokens(
-                raw_documents, batch_size=self.batch_size
-            )
-        flat_embeddings, lengths = flatten_repr(embeddings)
-        out_array = self.model.transform(
-            raw_documents, embeddings=flat_embeddings
-        )
+        chunks = get_document_chunks(raw_documents, offsets)
+        out_array = self.model.transform(chunks, embeddings=flat_embeddings)
         if self.pooling is None:
             return unflatten_repr(out_array, lengths)
         else:
@@ -141,14 +108,16 @@ def fit_transform(
         raw_documents: list[str],
         y=None,
         embeddings: list[np.ndarray] = None,
+        offsets: list[Offsets] = None,
     ):
-        if embeddings is None:
-            embeddings = self.model.encoder.encode_tokens(
+        if (embeddings is None) or (offsets is None):
+            embeddings, offsets = self.model.encoder.encode_tokens(
                 raw_documents, batch_size=self.batch_size
             )
         flat_embeddings, lengths = flatten_repr(embeddings)
+        chunks = get_document_chunks(raw_documents, offsets)
         out_array = self.model.fit_transform(
-            raw_documents, y, embeddings=flat_embeddings
+            chunks, embeddings=flat_embeddings
         )
         if self.pooling is None:
             return unflatten_repr(out_array, lengths)