feat: support batch processing with pruned_text for multiple texts

Previously, only the first raw_query/raw_text was sent to Python backend,
so process() was only called when batch_size == 1. Now all pairs are sent.

Changes:
- embed.proto: change to repeated string raw_queries/raw_texts
- grpc-client: accept Vec<String> instead of Option<String>
- backends/python/src/lib.rs: send all raw_queries/texts from batch
- types.py: extract lists from proto repeated fields
- xprovence_model.py: iterate batch and call process() for each pair

Now /rerank with multiple texts returns pruned_text for each result.

Author: sigridjineth

backends/grpc-client/src/client.rs

@@ -59,8 +59,8 @@ impl Client {
             position_ids,
             max_length,
             cu_seq_lengths,
-            raw_query: None,
-            raw_text: None,
+            raw_queries: vec![],
+            raw_texts: vec![],
         })
         .inject_context();
         let response = self.stub.embed(request).await?.into_inner();
@@ -75,17 +75,17 @@ impl Client {
         position_ids: Vec<u32>,
         cu_seq_lengths: Vec<u32>,
         max_length: u32,
-        raw_query: Option<String>,
-        raw_text: Option<String>,
+        raw_queries: Vec<String>,
+        raw_texts: Vec<String>,
     ) -> Result<Vec<Score>> {
         let request = tonic::Request::new(EmbedRequest {
             input_ids,
             token_type_ids,
             position_ids,
             max_length,
             cu_seq_lengths,
-            raw_query,
-            raw_text,
+            raw_queries,
+            raw_texts,
         })
         .inject_context();
         let response = self.stub.predict(request).await?.into_inner();

backends/proto/embed.proto

@@ -21,10 +21,10 @@ message EmbedRequest {
     repeated uint32 cu_seq_lengths = 4;
     /// Length of the longest request
     uint32 max_length = 5;
-    /// XProvence: raw query text for context pruning
-    optional string raw_query = 6;
-    /// XProvence: raw context text for context pruning
-    optional string raw_text = 7;
+    /// XProvence: raw query texts for context pruning (one per batch item)
+    repeated string raw_queries = 6;
+    /// XProvence: raw context texts for context pruning (one per batch item)
+    repeated string raw_texts = 7;
 }
 
 message Embedding {

backends/python/server/text_embeddings_server/models/types.py

@@ -3,7 +3,8 @@
 import torch
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import List, Optional
 from opentelemetry import trace
 
 from text_embeddings_server.pb import embed_pb2
@@ -36,9 +37,9 @@ class PaddedBatch(Batch):
     token_type_ids: torch.Tensor
     position_ids: torch.Tensor
     attention_mask: torch.Tensor
-    # XProvence: raw text for context pruning
-    raw_query: str = None
-    raw_text: str = None
+    # XProvence: raw texts for context pruning (one per batch item)
+    raw_queries: Optional[List[str]] = None
+    raw_texts: Optional[List[str]] = None
 
     @classmethod
     @tracer.start_as_current_span("from_pb")
@@ -80,27 +81,17 @@ def from_pb(
         # Move padded tensors all at once
         all_tensors = all_tensors.to(device)
 
-        # XProvence: Extract raw text if present in proto
-        # Use HasField for proto3 optional fields to properly detect if they were set
-        raw_query = None
-        raw_text = None
-        if hasattr(pb, 'HasField'):
-            if pb.HasField('raw_query'):
-                raw_query = pb.raw_query
-            if pb.HasField('raw_text'):
-                raw_text = pb.raw_text
-        else:
-            # Fallback for older proto versions
-            raw_query = pb.raw_query if pb.raw_query else None
-            raw_text = pb.raw_text if pb.raw_text else None
+        # XProvence: Extract repeated raw_queries/raw_texts from proto
+        raw_queries = list(pb.raw_queries) if pb.raw_queries else None
+        raw_texts = list(pb.raw_texts) if pb.raw_texts else None
 
         return PaddedBatch(
             input_ids=all_tensors[0],
             token_type_ids=all_tensors[1],
             position_ids=all_tensors[2],
             attention_mask=all_tensors[3],
-            raw_query=raw_query,
-            raw_text=raw_text,
+            raw_queries=raw_queries,
+            raw_texts=raw_texts,
         )
 
     def __len__(self):

backends/python/server/text_embeddings_server/models/xprovence_model.py

@@ -159,25 +159,45 @@ def predict(self, batch: PaddedBatch) -> List[Score]:
         """
         XProvence prediction with context pruning support.
 
-        For single-item batches with raw_query/raw_text available,
-        uses XProvence's process() method for sentence-level pruning.
+        For batches with raw_queries/raw_texts available (one per item),
+        uses XProvence's process() method for sentence-level pruning on each pair.
         Otherwise falls back to standard forward pass.
         """
         batch_size = len(batch)
 
-        # Debug: log raw_query/raw_text availability
-        has_query = batch.raw_query is not None
-        has_text = batch.raw_text is not None
+        # Check if we have raw data for the full batch
+        has_raw_data = (
+            batch.raw_queries is not None
+            and batch.raw_texts is not None
+            and len(batch.raw_queries) == batch_size
+            and len(batch.raw_texts) == batch_size
+        )
+
         logger.info(
             f"XProvence predict: batch_size={batch_size}, "
-            f"has_raw_query={has_query}, has_raw_text={has_text}"
+            f"has_raw_queries={batch.raw_queries is not None}, "
+            f"has_raw_texts={batch.raw_texts is not None}, "
+            f"has_full_raw_data={has_raw_data}"
         )
 
-        if batch_size == 1 and batch.raw_query and batch.raw_text:
-            logger.info("XProvence: Using process() for context pruning")
-            return self._predict_with_pruning(batch.raw_query, batch.raw_text)
-
-        logger.info("XProvence: Using standard forward pass (no raw_query/raw_text)")
+        if has_raw_data:
+            logger.info(f"XProvence: Processing batch of {batch_size} with pruning")
+            results = []
+            for i in range(batch_size):
+                query = batch.raw_queries[i]
+                text = batch.raw_texts[i]
+
+                # Verify we have valid strings (not empty)
+                if query and text:
+                    scores = self._predict_with_pruning(query, text)
+                    results.append(scores[0])
+                else:
+                    # Empty string fallback - use standard forward pass result
+                    logger.warning(f"XProvence: Empty query/text at index {i}, using 0.0")
+                    results.append(Score(values=[0.0], pruned_text=None))
+            return results
+
+        logger.info("XProvence: Using standard forward pass (no raw_queries/raw_texts)")
         return self._predict_standard(batch)
 
     def _predict_with_pruning(self, raw_query: str, raw_text: str) -> List[Score]:

backends/python/src/lib.rs

@@ -109,9 +109,17 @@ impl Backend for PythonBackend {
         }
         let batch_size = batch.len();
 
-        // XProvence: Get first raw query/text from batch (for single request)
-        let raw_query = batch.raw_queries.first().cloned().flatten();
-        let raw_text = batch.raw_texts.first().cloned().flatten();
+        // XProvence: Collect all raw queries/texts for the batch (one per item)
+        let raw_queries: Vec<String> = batch
+            .raw_queries
+            .into_iter()
+            .map(|q| q.unwrap_or_default())
+            .collect();
+        let raw_texts: Vec<String> = batch
+            .raw_texts
+            .into_iter()
+            .map(|t| t.unwrap_or_default())
+            .collect();
 
         let results = self
             .tokio_runtime
@@ -121,8 +129,8 @@ impl Backend for PythonBackend {
                 batch.position_ids,
                 batch.cumulative_seq_lengths,
                 batch.max_length,
-                raw_query,
-                raw_text,
+                raw_queries,
+                raw_texts,
             ))
             .map_err(|err| BackendError::Inference(err.to_string()))?;