dstpr 0.2.4__tar.gz → 0.2.5__tar.gz

{dstpr-0.2.4 → dstpr-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dstpr
-Version: 0.2.4
+Version: 0.2.5
 Summary: A high-throughput, domain-specific text preprocessing cascading pipeline filter to rank core sentences in texts and filter away boilerplate.
 Author-email: Stephen Meisenbacher <sjmeis@gtgd.com>
 Maintainer-email: Stephen Meisenbacher <sjmeis@gtgd.com>

{dstpr-0.2.4 → dstpr-0.2.5}/dstpr/pipeline.py

@@ -12,22 +12,38 @@ from dstpr.profiles import DomainProfile
 
 def get_optimal_gpu_batch_size(target_vram_ratio: float = 0.75) -> int:
     if not torch.cuda.is_available():
-        return 64
+        return 64, 32
     device = torch.cuda.current_device()
     free_bytes, _ = torch.cuda.mem_get_info(device)
     free_mb = free_bytes / (1024 ** 2)
-    # Estimate ~1.2MB memory per sample in FP16
-    ideal = int((free_mb * target_vram_ratio) / 1.2)
-    if ideal <= 128:
-        return 128
-    power_of_two = 1 << (ideal.bit_length() - 1)
-    return min(max(power_of_two, 128), 4096)
+    usable_mb = free_mb * target_vram_ratio
+
+    # Single-sentence Bi-Encoder: ~1.2 MB peak per sample (L <= 40)
+    ideal_bi = int(usable_mb / 1.2)
+    # Paired-input Cross-Encoder: ~4.5 MB peak per sample (L <= 128)
+    ideal_cross = int(usable_mb / 4.5)
+
+    def snap_pow2(val: int, min_val: int, max_val: int) -> int:
+        if val <= min_val:
+            return min_val
+        val = min(val, max_val)
+        return 1 << (val.bit_length() - 1)
+
+    bi_batch = snap_pow2(ideal_bi, min_val=128, max_val=2048)
+    cross_batch = snap_pow2(ideal_cross, min_val=64, max_val=512)
+
+    return bi_batch, cross_batch
 
 class ParallelPreprocessingPipeline:
-    def __init__(self, profile: DomainProfile, doc_chunk_size: int = 100, embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2", use_gpu: bool = False,):
+    def __init__(self, profile: DomainProfile, doc_chunk_size: int = 100, 
+                 embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2", 
+                 cross_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2", 
+                 use_gpu: bool = False
+        ):
         self.profile = profile
         self.doc_chunk_size = doc_chunk_size
         self.embedding_model = embedding_model
+        self.cross_model = cross_model
         self.use_gpu = use_gpu and torch.cuda.is_available()
         
         self.router = None
@@ -43,8 +59,6 @@ class ParallelPreprocessingPipeline:
             if self.use_gpu:
                 torch.backends.cuda.matmul.allow_tf32 = True
                 torch.backends.cudnn.allow_tf32 = True
-                self.inference_batch_size = get_optimal_gpu_batch_size()
-                print(f"GPU Auto-Calibrated Batch Size: {self.inference_batch_size}")
 
                 dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
                 self.bi_encoder = SentenceTransformer(
@@ -54,14 +68,16 @@ class ParallelPreprocessingPipeline:
                 )
 
                 self.cross_encoder = CrossEncoder(
-                    "cross-encoder/ms-marco-MiniLM-L6-v2",
+                    self.cross_model,
                     device="cuda",
                     model_kwargs={"dtype": dtype}
                 )
             else:
-                self.inference_batch_size = 256
                 self.bi_encoder = SentenceTransformer(self.embedding_model, device=device)
-                self.cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", device=device)
+                self.cross_encoder = CrossEncoder(self.cross_model, device=device)
+
+            self.bi_batch_size, self.cross_batch_size = get_optimal_gpu_batch_size()
+            print(f"Auto-Calibrated Batch Size: Bi | {self.bi_batch_size}, Cross | {self.cross_batch_size}")
 
             self.router = SemanticSectionRouter(profile=self.profile, model=self.bi_encoder)
             self.ranker = HybridRanker(
@@ -81,7 +97,6 @@ class ParallelPreprocessingPipeline:
                 continue
 
             sents = clean_and_split_chunks(raw_text, self.profile)
-            n_sents = len(sents)
             doc_sentence_counts.append(len(sents))
             flat_sentences.extend(sents)
 
@@ -102,7 +117,7 @@ class ParallelPreprocessingPipeline:
         flat_candidates, flat_cand_embs, doc_cand_counts = self.router.route_sentences(
             sentences=flat_sentences,
             doc_sentence_counts=doc_sentence_counts,
-            batch_size=self.inference_batch_size,
+            batch_size=self.bi_batch_size
         )
 
         filtered_doc_sentences = self.ranker.rank_and_filter(
@@ -110,7 +125,7 @@ class ParallelPreprocessingPipeline:
             flat_cand_embs=flat_cand_embs,
             doc_cand_counts=doc_cand_counts,
             cross_prob_threshold=threshold,
-            batch_size=self.inference_batch_size,
+            batch_size=self.cross_batch_size
         )
 
         return [sanitize_outputs(sents, self.profile) for sents in filtered_doc_sentences]

{dstpr-0.2.4 → dstpr-0.2.5}/dstpr.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dstpr
-Version: 0.2.4
+Version: 0.2.5
 Summary: A high-throughput, domain-specific text preprocessing cascading pipeline filter to rank core sentences in texts and filter away boilerplate.
 Author-email: Stephen Meisenbacher <sjmeis@gtgd.com>
 Maintainer-email: Stephen Meisenbacher <sjmeis@gtgd.com>

{dstpr-0.2.4 → dstpr-0.2.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dstpr"
-version = "0.2.4"
+version = "0.2.5"
 description = "A high-throughput, domain-specific text preprocessing cascading pipeline filter to rank core sentences in texts and filter away boilerplate."
 readme = {file = "README.md", content-type = "text/markdown"}
 license = {file = "LICENSE"}