1#[derive(Debug, Clone)]
13pub struct ExtractorSummaryConfig {
14 pub max_sentences: usize,
16 pub similarity_threshold: f64,
18 pub diversity_penalty: f64,
21}
22
23impl Default for ExtractorSummaryConfig {
24 fn default() -> Self {
25 Self {
26 max_sentences: 5,
27 similarity_threshold: 0.3,
28 diversity_penalty: 0.5,
29 }
30 }
31}
32
33#[derive(Debug, Clone)]
37pub struct ExtractorScoredSentence {
38 pub index: usize,
40 pub text: String,
42 pub embedding: Vec<f64>,
44 pub score: f64,
46 pub selected: bool,
48}
49
50#[derive(Debug, Clone)]
52pub struct ExtractionResult {
53 pub selected_indices: Vec<usize>,
55 pub sentences: Vec<ExtractorScoredSentence>,
57 pub coverage_score: f64,
60}
61
62#[derive(Debug, Clone)]
64pub struct SummaryExtractorStats {
65 pub extractions_performed: u64,
67}
68
69pub struct SemanticSummaryExtractor {
73 config: ExtractorSummaryConfig,
74 extractions_performed: u64,
75}
76
77impl SemanticSummaryExtractor {
78 pub fn new(config: ExtractorSummaryConfig) -> Self {
80 Self {
81 config,
82 extractions_performed: 0,
83 }
84 }
85
86 pub fn extract(
92 &mut self,
93 sentences: &[(String, Vec<f64>)],
94 query_embedding: Option<&[f64]>,
95 ) -> Result<ExtractionResult, String> {
96 if sentences.is_empty() {
97 return Err("input sentences must not be empty".to_string());
98 }
99
100 let embeddings: Vec<&Vec<f64>> = sentences.iter().map(|(_, e)| e).collect();
101
102 let base_scores: Vec<f64> = if let Some(query) = query_embedding {
104 embeddings
105 .iter()
106 .map(|e| Self::cosine_similarity(e, query))
107 .collect()
108 } else {
109 let embs: Vec<Vec<f64>> = embeddings.iter().map(|e| (*e).clone()).collect();
110 Self::centrality_scores(&embs)
111 };
112
113 let mut scores = base_scores.clone();
115 let mut selected_flags = vec![false; sentences.len()];
116 let mut selected_indices: Vec<usize> = Vec::new();
117
118 for _ in 0..self.config.max_sentences {
120 let best = scores
122 .iter()
123 .enumerate()
124 .filter(|(i, _)| !selected_flags[*i])
125 .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
126
127 let (best_idx, &best_score) = match best {
128 Some(pair) => pair,
129 None => break, };
131
132 if best_score < self.config.similarity_threshold {
133 break;
134 }
135
136 selected_flags[best_idx] = true;
137 selected_indices.push(best_idx);
138
139 for (j, s) in scores.iter_mut().enumerate() {
141 if !selected_flags[j] {
142 let sim = Self::cosine_similarity(embeddings[j], embeddings[best_idx]);
143 *s -= self.config.diversity_penalty * sim;
144 }
145 }
146 }
147
148 let scored: Vec<ExtractorScoredSentence> = sentences
150 .iter()
151 .enumerate()
152 .map(|(i, (text, emb))| ExtractorScoredSentence {
153 index: i,
154 text: text.clone(),
155 embedding: emb.clone(),
156 score: scores[i],
157 selected: selected_flags[i],
158 })
159 .collect();
160
161 let selected_embs: Vec<Vec<f64>> = selected_indices
162 .iter()
163 .map(|&i| sentences[i].1.clone())
164 .collect();
165 let all_embs: Vec<Vec<f64>> = sentences.iter().map(|(_, e)| e.clone()).collect();
166 let coverage_score = Self::coverage(&selected_embs, &all_embs);
167
168 self.extractions_performed += 1;
169
170 Ok(ExtractionResult {
171 selected_indices,
172 sentences: scored,
173 coverage_score,
174 })
175 }
176
177 pub fn cosine_similarity(a: &[f64], b: &[f64]) -> f64 {
183 let dot: f64 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
184 let mag_a: f64 = a.iter().map(|x| x * x).sum::<f64>().sqrt();
185 let mag_b: f64 = b.iter().map(|x| x * x).sum::<f64>().sqrt();
186 if mag_a == 0.0 || mag_b == 0.0 {
187 return 0.0;
188 }
189 dot / (mag_a * mag_b)
190 }
191
192 pub fn centrality_scores(embeddings: &[Vec<f64>]) -> Vec<f64> {
194 let n = embeddings.len();
195 if n <= 1 {
196 return vec![0.0; n];
197 }
198 let mut scores = vec![0.0_f64; n];
199 for i in 0..n {
200 for j in 0..n {
201 if i != j {
202 scores[i] += Self::cosine_similarity(&embeddings[i], &embeddings[j]);
203 }
204 }
205 scores[i] /= (n - 1) as f64;
206 }
207 scores
208 }
209
210 pub fn coverage(selected: &[Vec<f64>], all: &[Vec<f64>]) -> f64 {
216 if selected.is_empty() || all.is_empty() {
217 return 0.0;
218 }
219 let total: f64 = all
220 .iter()
221 .map(|a| {
222 selected
223 .iter()
224 .map(|s| Self::cosine_similarity(a, s))
225 .fold(f64::NEG_INFINITY, f64::max)
226 })
227 .sum();
228 total / all.len() as f64
229 }
230
231 pub fn stats(&self) -> SummaryExtractorStats {
233 SummaryExtractorStats {
234 extractions_performed: self.extractions_performed,
235 }
236 }
237}
238
239#[cfg(test)]
242mod tests {
243 use super::*;
244
245 fn default_extractor() -> SemanticSummaryExtractor {
246 SemanticSummaryExtractor::new(ExtractorSummaryConfig::default())
247 }
248
249 fn make_sentences(vecs: &[Vec<f64>]) -> Vec<(String, Vec<f64>)> {
250 vecs.iter()
251 .enumerate()
252 .map(|(i, v)| (format!("sentence {i}"), v.clone()))
253 .collect()
254 }
255
256 #[test]
259 fn cosine_parallel_vectors() {
260 let a = vec![1.0, 0.0, 0.0];
261 let b = vec![2.0, 0.0, 0.0];
262 let sim = SemanticSummaryExtractor::cosine_similarity(&a, &b);
263 assert!(
264 (sim - 1.0).abs() < 1e-9,
265 "parallel vectors should have similarity 1.0"
266 );
267 }
268
269 #[test]
270 fn cosine_orthogonal_vectors() {
271 let a = vec![1.0, 0.0];
272 let b = vec![0.0, 1.0];
273 let sim = SemanticSummaryExtractor::cosine_similarity(&a, &b);
274 assert!(
275 sim.abs() < 1e-9,
276 "orthogonal vectors should have similarity 0.0"
277 );
278 }
279
280 #[test]
281 fn cosine_antiparallel_vectors() {
282 let a = vec![1.0, 0.0];
283 let b = vec![-1.0, 0.0];
284 let sim = SemanticSummaryExtractor::cosine_similarity(&a, &b);
285 assert!(
286 (sim + 1.0).abs() < 1e-9,
287 "antiparallel vectors should have similarity -1.0"
288 );
289 }
290
291 #[test]
292 fn cosine_zero_vector() {
293 let a = vec![0.0, 0.0];
294 let b = vec![1.0, 2.0];
295 assert_eq!(SemanticSummaryExtractor::cosine_similarity(&a, &b), 0.0);
296 }
297
298 #[test]
299 fn cosine_identical_vectors() {
300 let a = vec![0.3, 0.4, 0.5];
301 let sim = SemanticSummaryExtractor::cosine_similarity(&a, &a);
302 assert!((sim - 1.0).abs() < 1e-9);
303 }
304
305 #[test]
308 fn centrality_single_embedding() {
309 let embs = vec![vec![1.0, 0.0]];
310 let scores = SemanticSummaryExtractor::centrality_scores(&embs);
311 assert_eq!(scores, vec![0.0]);
312 }
313
314 #[test]
315 fn centrality_two_identical() {
316 let embs = vec![vec![1.0, 0.0], vec![1.0, 0.0]];
317 let scores = SemanticSummaryExtractor::centrality_scores(&embs);
318 assert!((scores[0] - 1.0).abs() < 1e-9);
319 assert!((scores[1] - 1.0).abs() < 1e-9);
320 }
321
322 #[test]
323 fn centrality_orthogonal_pair() {
324 let embs = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
325 let scores = SemanticSummaryExtractor::centrality_scores(&embs);
326 assert!(scores[0].abs() < 1e-9);
327 assert!(scores[1].abs() < 1e-9);
328 }
329
330 #[test]
331 fn centrality_three_embeddings() {
332 let embs = vec![vec![1.0, 0.0], vec![1.0, 0.0], vec![0.0, 1.0]];
333 let scores = SemanticSummaryExtractor::centrality_scores(&embs);
334 assert!((scores[0] - 0.5).abs() < 1e-9);
336 assert!((scores[1] - 0.5).abs() < 1e-9);
337 assert!((scores[2] - 0.0).abs() < 1e-9);
338 }
339
340 #[test]
343 fn extract_with_query_selects_most_similar() {
344 let mut ext = default_extractor();
345 let sents = make_sentences(&[vec![1.0, 0.0], vec![0.0, 1.0], vec![0.7, 0.7]]);
346 let query = vec![1.0, 0.0];
347 let res = ext.extract(&sents, Some(&query)).expect("should succeed");
348 assert_eq!(res.selected_indices[0], 0);
350 }
351
352 #[test]
353 fn extract_with_query_respects_max_sentences() {
354 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
355 max_sentences: 2,
356 similarity_threshold: 0.0,
357 diversity_penalty: 0.0,
358 });
359 let sents = make_sentences(&[
360 vec![1.0, 0.0],
361 vec![0.9, 0.1],
362 vec![0.8, 0.2],
363 vec![0.7, 0.3],
364 ]);
365 let query = vec![1.0, 0.0];
366 let res = ext.extract(&sents, Some(&query)).expect("should succeed");
367 assert_eq!(res.selected_indices.len(), 2);
368 }
369
370 #[test]
371 fn extract_with_query_all_selected() {
372 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
373 max_sentences: 10,
374 similarity_threshold: 0.0,
375 diversity_penalty: 0.0,
376 });
377 let sents = make_sentences(&[vec![1.0, 0.0], vec![0.0, 1.0]]);
378 let query = vec![1.0, 1.0];
379 let res = ext.extract(&sents, Some(&query)).expect("should succeed");
380 assert_eq!(res.selected_indices.len(), 2);
381 }
382
383 #[test]
386 fn extract_without_query_uses_centrality() {
387 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
388 max_sentences: 1,
389 similarity_threshold: 0.0,
390 diversity_penalty: 0.0,
391 });
392 let sents = make_sentences(&[vec![1.0, 0.1], vec![1.0, 0.0], vec![0.0, 1.0]]);
395 let res = ext.extract(&sents, None).expect("should succeed");
396 assert!(
397 res.selected_indices[0] == 0 || res.selected_indices[0] == 1,
398 "should select one of the two similar sentences"
399 );
400 }
401
402 #[test]
403 fn extract_centrality_with_diversity() {
404 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
405 max_sentences: 2,
406 similarity_threshold: -10.0,
407 diversity_penalty: 0.5,
408 });
409 let sents = make_sentences(&[
410 vec![1.0, 0.0],
411 vec![1.0, 0.01], vec![0.0, 1.0], ]);
414 let res = ext.extract(&sents, None).expect("should succeed");
415 assert_eq!(res.selected_indices.len(), 2);
418 assert_ne!(res.selected_indices[0], res.selected_indices[1]);
420 }
421
422 #[test]
425 fn diversity_penalty_reduces_redundancy() {
426 let sents = make_sentences(&[vec![1.0, 0.0], vec![0.99, 0.01], vec![0.0, 1.0]]);
429 let query = vec![1.0, 0.0];
430
431 let mut no_penalty = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
432 max_sentences: 2,
433 similarity_threshold: 0.0,
434 diversity_penalty: 0.0,
435 });
436 let r1 = no_penalty.extract(&sents, Some(&query)).expect("ok");
437 assert_eq!(r1.selected_indices, vec![0, 1]);
439
440 let mut with_penalty = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
441 max_sentences: 2,
442 similarity_threshold: 0.0,
443 diversity_penalty: 2.0, });
445 let r2 = with_penalty.extract(&sents, Some(&query)).expect("ok");
446 assert_eq!(r2.selected_indices[0], 0);
448 assert_eq!(r2.selected_indices[1], 2);
449 }
450
451 #[test]
454 fn max_sentences_caps_output() {
455 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
456 max_sentences: 1,
457 similarity_threshold: 0.0,
458 diversity_penalty: 0.0,
459 });
460 let sents = make_sentences(&[vec![1.0, 0.0], vec![0.0, 1.0], vec![0.5, 0.5]]);
461 let res = ext.extract(&sents, Some(&[1.0, 0.0])).expect("ok");
462 assert_eq!(res.selected_indices.len(), 1);
463 }
464
465 #[test]
468 fn empty_input_returns_error() {
469 let mut ext = default_extractor();
470 let res = ext.extract(&[], None);
471 assert!(res.is_err());
472 }
473
474 #[test]
477 fn single_sentence_selected() {
478 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
479 max_sentences: 5,
480 similarity_threshold: 0.0,
481 diversity_penalty: 0.5,
482 });
483 let sents = make_sentences(&[vec![1.0, 0.0]]);
484 let res = ext.extract(&sents, None).expect("ok");
486 assert_eq!(res.selected_indices.len(), 1);
487 assert_eq!(res.selected_indices[0], 0);
488 }
489
490 #[test]
491 fn single_sentence_with_query() {
492 let mut ext = default_extractor();
493 let sents = make_sentences(&[vec![1.0, 0.0]]);
494 let res = ext.extract(&sents, Some(&[1.0, 0.0])).expect("ok");
495 assert_eq!(res.selected_indices.len(), 1);
496 }
497
498 #[test]
501 fn all_below_threshold() {
502 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
503 max_sentences: 5,
504 similarity_threshold: 0.99,
505 diversity_penalty: 0.0,
506 });
507 let sents = make_sentences(&[vec![0.0, 1.0]]);
509 let res = ext.extract(&sents, Some(&[1.0, 0.0])).expect("ok");
510 assert!(res.selected_indices.is_empty());
511 }
512
513 #[test]
514 fn threshold_filters_partial() {
515 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
516 max_sentences: 10,
517 similarity_threshold: 0.9,
518 diversity_penalty: 0.0,
519 });
520 let sents = make_sentences(&[
521 vec![1.0, 0.0], vec![0.0, 1.0], ]);
524 let res = ext.extract(&sents, Some(&[1.0, 0.0])).expect("ok");
525 assert_eq!(res.selected_indices, vec![0]);
526 }
527
528 #[test]
531 fn coverage_perfect_when_all_selected() {
532 let all = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
533 let cov = SemanticSummaryExtractor::coverage(&all, &all);
534 assert!(
535 (cov - 1.0).abs() < 1e-9,
536 "coverage should be 1.0 when all selected"
537 );
538 }
539
540 #[test]
541 fn coverage_zero_when_none_selected() {
542 let all = vec![vec![1.0, 0.0]];
543 let cov = SemanticSummaryExtractor::coverage(&[], &all);
544 assert_eq!(cov, 0.0);
545 }
546
547 #[test]
548 fn coverage_partial() {
549 let selected = vec![vec![1.0, 0.0]];
550 let all = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
551 let cov = SemanticSummaryExtractor::coverage(&selected, &all);
552 assert!((cov - 0.5).abs() < 1e-9);
556 }
557
558 #[test]
559 fn coverage_with_similar_sentences() {
560 let selected = vec![vec![1.0, 0.0]];
561 let all = vec![vec![1.0, 0.0], vec![0.9, 0.1]];
562 let cov = SemanticSummaryExtractor::coverage(&selected, &all);
563 assert!(cov > 0.99);
567 }
568
569 #[test]
572 fn stats_tracks_extractions() {
573 let mut ext = default_extractor();
574 assert_eq!(ext.stats().extractions_performed, 0);
575
576 let sents = make_sentences(&[vec![1.0, 0.0]]);
577 let _ = ext.extract(&sents, Some(&[1.0, 0.0]));
578 assert_eq!(ext.stats().extractions_performed, 1);
579
580 let _ = ext.extract(&sents, None);
581 assert_eq!(ext.stats().extractions_performed, 2);
582 }
583
584 #[test]
585 fn stats_not_incremented_on_error() {
586 let mut ext = default_extractor();
587 let _ = ext.extract(&[], None); assert_eq!(ext.stats().extractions_performed, 0);
589 }
590
591 #[test]
594 fn deterministic_results() {
595 let sents = make_sentences(&[vec![1.0, 0.0], vec![0.5, 0.5], vec![0.0, 1.0]]);
596 let query = vec![0.6, 0.4];
597
598 let mut ext1 = default_extractor();
599 let mut ext2 = default_extractor();
600
601 let r1 = ext1.extract(&sents, Some(&query)).expect("ok");
602 let r2 = ext2.extract(&sents, Some(&query)).expect("ok");
603
604 assert_eq!(r1.selected_indices, r2.selected_indices);
605 assert!((r1.coverage_score - r2.coverage_score).abs() < 1e-12);
606 }
607
608 #[test]
611 fn selected_flags_match_indices() {
612 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
613 max_sentences: 2,
614 similarity_threshold: 0.0,
615 diversity_penalty: 0.0,
616 });
617 let sents = make_sentences(&[vec![1.0, 0.0], vec![0.5, 0.5], vec![0.0, 1.0]]);
618 let res = ext.extract(&sents, Some(&[1.0, 0.0])).expect("ok");
619
620 for sent in &res.sentences {
621 if res.selected_indices.contains(&sent.index) {
622 assert!(sent.selected);
623 } else {
624 assert!(!sent.selected);
625 }
626 }
627 }
628
629 #[test]
632 fn default_config_values() {
633 let cfg = ExtractorSummaryConfig::default();
634 assert_eq!(cfg.max_sentences, 5);
635 assert!((cfg.similarity_threshold - 0.3).abs() < 1e-9);
636 assert!((cfg.diversity_penalty - 0.5).abs() < 1e-9);
637 }
638
639 #[test]
642 fn high_dimensional_embeddings() {
643 let dim = 128;
644 let mut ext = default_extractor();
645 let mut v1 = vec![0.0; dim];
646 v1[0] = 1.0;
647 let mut v2 = vec![0.0; dim];
648 v2[1] = 1.0;
649 let mut v3 = vec![0.0; dim];
650 v3[0] = 0.7;
651 v3[1] = 0.7;
652
653 let sents = make_sentences(&[v1.clone(), v2, v3]);
654 let res = ext.extract(&sents, Some(&v1)).expect("ok");
655 assert!(!res.selected_indices.is_empty());
656 }
657
658 #[test]
661 fn extract_result_coverage_is_consistent() {
662 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
663 max_sentences: 2,
664 similarity_threshold: 0.0,
665 diversity_penalty: 0.0,
666 });
667 let sents = make_sentences(&[vec![1.0, 0.0], vec![0.0, 1.0], vec![0.5, 0.5]]);
668 let res = ext.extract(&sents, Some(&[1.0, 0.0])).expect("ok");
669
670 let selected_embs: Vec<Vec<f64>> = res
672 .selected_indices
673 .iter()
674 .map(|&i| sents[i].1.clone())
675 .collect();
676 let all_embs: Vec<Vec<f64>> = sents.iter().map(|(_, e)| e.clone()).collect();
677 let expected_cov = SemanticSummaryExtractor::coverage(&selected_embs, &all_embs);
678
679 assert!((res.coverage_score - expected_cov).abs() < 1e-12);
680 }
681
682 #[test]
685 fn all_identical_sentences() {
686 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
687 max_sentences: 3,
688 similarity_threshold: 0.0,
689 diversity_penalty: 0.5,
690 });
691 let sents = make_sentences(&[vec![1.0, 0.0], vec![1.0, 0.0], vec![1.0, 0.0]]);
692 let res = ext.extract(&sents, Some(&[1.0, 0.0])).expect("ok");
693 assert!(!res.selected_indices.is_empty());
695 }
696
697 #[test]
700 fn negative_scores_below_threshold_not_selected() {
701 let mut ext = SemanticSummaryExtractor::new(ExtractorSummaryConfig {
702 max_sentences: 5,
703 similarity_threshold: 0.3,
704 diversity_penalty: 5.0, });
706 let sents = make_sentences(&[vec![1.0, 0.0], vec![0.9, 0.1], vec![0.8, 0.2]]);
707 let res = ext.extract(&sents, Some(&[1.0, 0.0])).expect("ok");
708 assert!(res.selected_indices.contains(&0));
711 assert!(res.selected_indices.len() <= 2);
713 }
714}