Update README.md

README.md

@@ -96,8 +96,68 @@ def reform(features, how='softmax'):
         return torch.nn.functional.normalize(features.reshape(1,-1))
 
 feature_vector = reform(get_feature_vector(text))
+```
+
+or with multiple texts
+
+```python
+
+def get_feature_vectors_with_stride(
+    texts,
+    stride=16,
+    max_context_length=128,
+    reformer=None
+):
+    """
+    Extract one feature vector per document.
+
+    Long documents are split into overlapping windows.
+    Logits from all windows belonging to the same document are averaged.
+
+    Returns:
+        Tensor [n_documents, n_features]
+    """
+
+    inputs = tokenizer(
+        texts,
+        return_tensors="pt",
+        truncation=True,
+        max_length=max_context_length,
+        stride=stride,
+        return_overflowing_tokens=True,
+        padding=True,
+    )
+
+    # Maps each generated chunk back to its original document.
+    sample_mapping = inputs.pop("overflow_to_sample_mapping")
+
+    inputs = {
+        key: value.to(device)
+        for key, value in inputs.items()
+    }
+
+    with torch.no_grad():
+        logits = model(**inputs).logits
+
+    # logits:
+    # [n_chunks, n_features]
+
+    sample_mapping = sample_mapping.to(logits.device)
+
+    features = []
+
+    for sample_idx in range(len(texts)):
+        mask = sample_mapping == sample_idx
+
+        # Mean over all windows belonging to this document.
+        if reformer is None:
+            document_features = logits[mask].mean(dim=0)
+        else:
+            document_features = reform(logits[mask].mean(dim=0).flatten(), how=reformer)
 
+        features.append(document_features)
 
+    return torch.stack(features).cpu()
 ```