Method encode

rag_system/indexing/latechunk.py:43–88

Return one vector *per* span. Args: text: Full document text. chunk_spans: List of (char_start, char_end) offsets for each chunk. Returns: List of numpy float32 arrays – one per chunk.

(self, text: str, chunk_spans: List[Tuple[int, int]])

    @torch.inference_mode()
    def encode(self, text: str, chunk_spans: List[Tuple[int, int]]) -> List[np.ndarray]:
        """Return one vector *per* span.

        Args:
            text: Full document text.
            chunk_spans: List of (char_start, char_end) offsets for each chunk.

        Returns:
            List of numpy float32 arrays – one per chunk.
        """
        if not chunk_spans:
            return []

        # Tokenise and obtain per-token hidden states
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            return_offsets_mapping=True,
            truncation=True,
            max_length=self.max_len,
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        offsets = inputs.pop("offset_mapping").squeeze(0).cpu().tolist()  # (seq_len, 2)

        out = self.model(**inputs)
        last_hidden = out.last_hidden_state.squeeze(0)  # (seq_len, dim)
        last_hidden = last_hidden.cpu()

        # For each chunk span, gather token indices belonging to it
        vectors: List[np.ndarray] = []
        for start_char, end_char in chunk_spans:
            token_indices = [i for i, (s, e) in enumerate(offsets) if s >= start_char and e <= end_char]
            if not token_indices:
                # Fallback: if tokenizer lost the span (e.g. due to trimming) just average CLS + SEP
                token_indices = [0]
            chunk_vec = last_hidden[token_indices].mean(dim=0).numpy().astype("float32")
            
            # Check for NaN or infinite values
            if np.isnan(chunk_vec).any() or np.isinf(chunk_vec).any():
                print(f"⚠️ Warning: Invalid values detected in late chunk embedding for span ({start_char}, {end_char})")
                # Replace invalid values with zeros
                chunk_vec = np.nan_to_num(chunk_vec, nan=0.0, posinf=0.0, neginf=0.0)
                print(f"🔄 Replaced invalid values with zeros")
            
            vectors.append(chunk_vec)
        return vectors