anno_eval/eval/
cross_context_eval.rs

1//! Cross-Context Coreference Evaluation Harness.
2//!
3//! Evaluation framework for xCoRe-style cross-context coreference resolution,
4//! supporting both long-document and cross-document benchmarks.
5//!
6//! # Supported Benchmarks
7//!
8//! This harness is intended for ECB+, SciCo, LitBank, BookCoref, and similar
9//! long-document / cross-document benchmarks.
10//!
11//! # Usage
12//!
13//! ```rust,ignore
14//! use anno_eval::eval::cross_context_eval::{CrossContextBenchmark, evaluate_benchmark};
15//! use anno::metrics::cluster_encoder::{HeuristicClusterEncoder, CosineMergeScorer};
16//!
17//! let encoder = HeuristicClusterEncoder::new(64);
18//! let scorer = CosineMergeScorer::new();
19//!
20//! let results = evaluate_benchmark(
21//!     CrossContextBenchmark::ECBPlus,
22//!     &encoder,
23//!     &scorer,
24//!     None, // Use default config
25//! )?;
26//!
27//! println!("CoNLL F1: {:.1}", results.conll_f1 * 100.0);
28//! ```
29//!
30//! # References
31//!
32//! - Martinelli et al. (2025): "xCoRe: Cross-context Coreference Resolution"
33//! - Cybulska & Vossen (2014): "ECB+ Event Coreference Bank"
34//! - Cattan et al. (2021): "SciCo Hierarchical Cross-Document Coreference"
35//! - Bamman et al. (2020): "LitBank"
36//! - Martinelli et al. (2025): "BOOKCOREF: Coreference Resolution at Book Scale"
37//! - Guo et al. (2023): "Animal Farm annotation"
38
39use crate::eval::cdcr::{CrossDocCluster, Document};
40use crate::eval::cluster_encoder::{ClusterEncoder, MergeScorer};
41use crate::eval::coref::{CorefChain, Mention};
42use crate::eval::coref_metrics::{conll_f1, CorefScores};
43use crate::eval::neural_cluster_encoder::{
44    CrossContextConfig, UnifiedCrossContextResolver, WindowOutput,
45};
46use crate::Result;
47use serde::{Deserialize, Serialize};
48use std::collections::HashMap;
49
50// =============================================================================
51// Benchmark Definitions
52// =============================================================================
53
54/// Cross-context coreference benchmarks.
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
56pub enum CrossContextBenchmark {
57    /// ECB+ - Cross-document entity/event coreference (news)
58    ECBPlus,
59    /// SciCo - Cross-document concept coreference (scientific papers)
60    SciCo,
61    /// LitBank - Long-document coreference (literary fiction)
62    LitBank,
63    /// BookCoref - Full-book coreference (book-scale)
64    BookCoref,
65    /// Animal Farm - Single long novel benchmark
66    AnimalFarm,
67}
68
69impl CrossContextBenchmark {
70    /// Get benchmark name.
71    pub fn name(&self) -> &'static str {
72        match self {
73            Self::ECBPlus => "ECB+",
74            Self::SciCo => "SciCo",
75            Self::LitBank => "LitBank",
76            Self::BookCoref => "BookCoref",
77            Self::AnimalFarm => "Animal Farm",
78        }
79    }
80
81    /// Is this a cross-document benchmark?
82    pub fn is_cross_document(&self) -> bool {
83        matches!(self, Self::ECBPlus | Self::SciCo)
84    }
85
86    /// Is this a long-document benchmark?
87    pub fn is_long_document(&self) -> bool {
88        matches!(self, Self::LitBank | Self::BookCoref | Self::AnimalFarm)
89    }
90
91    /// Get recommended window size for this benchmark.
92    pub fn recommended_window_size(&self) -> usize {
93        match self {
94            Self::ECBPlus => 512,     // Documents are short
95            Self::SciCo => 512,       // Paper sections
96            Self::LitBank => 2000,    // Limited to 2k tokens
97            Self::BookCoref => 4000,  // Full books
98            Self::AnimalFarm => 4000, // Single long novel
99        }
100    }
101
102    /// State-of-the-art CoNLL F1 from xCoRe (Martinelli et al. 2025, Table 3).
103    pub fn xcore_sota_f1(&self) -> f64 {
104        match self {
105            Self::ECBPlus => 40.3,
106            Self::SciCo => 34.5,
107            Self::LitBank => 78.2,
108            Self::BookCoref => 65.0,
109            Self::AnimalFarm => 70.0,
110        }
111    }
112
113    /// Get all benchmarks.
114    pub fn all() -> &'static [Self] {
115        &[
116            Self::ECBPlus,
117            Self::SciCo,
118            Self::LitBank,
119            Self::BookCoref,
120            Self::AnimalFarm,
121        ]
122    }
123
124    /// Get cross-document benchmarks.
125    pub fn cross_document() -> &'static [Self] {
126        &[Self::ECBPlus, Self::SciCo]
127    }
128
129    /// Get long-document benchmarks.
130    pub fn long_document() -> &'static [Self] {
131        &[Self::LitBank, Self::BookCoref, Self::AnimalFarm]
132    }
133}
134
135// =============================================================================
136// Evaluation Configuration
137// =============================================================================
138
139/// Configuration for cross-context evaluation.
140#[derive(Debug, Clone, Serialize, Deserialize)]
141pub struct CrossContextEvalConfig {
142    /// Window size for long-document processing
143    pub window_size: usize,
144    /// Window overlap
145    pub window_overlap: usize,
146    /// Merge probability threshold
147    pub merge_threshold: f32,
148    /// Whether to use gold mentions (vs predicted)
149    pub use_gold_mentions: bool,
150    /// Whether to use gold within-context clusters
151    pub use_gold_clusters: bool,
152    /// Maximum documents per topic (for cross-doc, 0 = all)
153    pub max_docs_per_topic: usize,
154    /// Random seed for sampling
155    pub seed: u64,
156}
157
158impl Default for CrossContextEvalConfig {
159    fn default() -> Self {
160        Self {
161            window_size: 4000,
162            window_overlap: 256,
163            merge_threshold: 0.5,
164            use_gold_mentions: false,
165            use_gold_clusters: false,
166            max_docs_per_topic: 0,
167            seed: 42,
168        }
169    }
170}
171
172impl CrossContextEvalConfig {
173    /// Create config for a specific benchmark.
174    pub fn for_benchmark(benchmark: CrossContextBenchmark) -> Self {
175        Self {
176            window_size: benchmark.recommended_window_size(),
177            ..Default::default()
178        }
179    }
180
181    /// Config for oracle evaluation (gold mentions + gold clusters).
182    pub fn oracle() -> Self {
183        Self {
184            use_gold_mentions: true,
185            use_gold_clusters: true,
186            ..Default::default()
187        }
188    }
189
190    /// Config for predicted mentions, gold clusters.
191    pub fn gold_clusters() -> Self {
192        Self {
193            use_gold_mentions: false,
194            use_gold_clusters: true,
195            ..Default::default()
196        }
197    }
198}
199
200// =============================================================================
201// Evaluation Results
202// =============================================================================
203
204/// Results from cross-context evaluation.
205#[derive(Debug, Clone, Serialize, Deserialize)]
206pub struct CrossContextEvalResults {
207    /// Benchmark name
208    pub benchmark: String,
209    /// Configuration used
210    pub config: CrossContextEvalConfig,
211    /// MUC scores
212    pub muc: CorefScores,
213    /// B³ scores
214    pub b_cubed: CorefScores,
215    /// CEAF-e scores
216    pub ceaf_e: CorefScores,
217    /// LEA scores
218    pub lea: CorefScores,
219    /// CoNLL F1 (average of MUC, B³, CEAF-e)
220    pub conll_f1: f64,
221    /// Number of documents/contexts evaluated
222    pub num_contexts: usize,
223    /// Number of gold clusters
224    pub num_gold_clusters: usize,
225    /// Number of predicted clusters
226    pub num_pred_clusters: usize,
227    /// Average cluster size
228    pub avg_cluster_size: f64,
229    /// Processing time in milliseconds
230    pub time_ms: f64,
231    /// Per-topic results (for cross-document)
232    pub per_topic: Option<HashMap<String, TopicResults>>,
233    /// Per-document results (for long-document)
234    pub per_document: Option<HashMap<String, DocumentResults>>,
235}
236
237impl CrossContextEvalResults {
238    /// Format as summary string.
239    pub fn summary(&self) -> String {
240        format!(
241            "{}: CoNLL F1 = {:.1}% (MUC: {:.1}, B³: {:.1}, CEAF: {:.1})",
242            self.benchmark,
243            self.conll_f1 * 100.0,
244            self.muc.f1 * 100.0,
245            self.b_cubed.f1 * 100.0,
246            self.ceaf_e.f1 * 100.0,
247        )
248    }
249}
250
251/// Per-topic results for cross-document evaluation.
252#[derive(Debug, Clone, Serialize, Deserialize)]
253pub struct TopicResults {
254    /// Topic ID
255    pub topic_id: String,
256    /// Number of documents in topic
257    pub num_documents: usize,
258    /// CoNLL F1 for this topic
259    pub conll_f1: f64,
260    /// Number of gold clusters
261    pub num_gold_clusters: usize,
262    /// Number of predicted clusters
263    pub num_pred_clusters: usize,
264}
265
266/// Per-document results for long-document evaluation.
267#[derive(Debug, Clone, Serialize, Deserialize)]
268pub struct DocumentResults {
269    /// Document ID
270    pub doc_id: String,
271    /// Document length in tokens
272    pub num_tokens: usize,
273    /// Number of windows
274    pub num_windows: usize,
275    /// CoNLL F1 for this document
276    pub conll_f1: f64,
277    /// Number of gold chains
278    pub num_gold_chains: usize,
279    /// Number of predicted chains
280    pub num_pred_chains: usize,
281}
282
283// =============================================================================
284// Example Data Structures
285// =============================================================================
286
287/// A topic containing multiple documents (for cross-document evaluation).
288#[derive(Debug, Clone)]
289pub struct Topic {
290    /// Topic ID
291    pub id: String,
292    /// Documents in this topic
293    pub documents: Vec<Document>,
294    /// Gold cross-document clusters
295    pub gold_clusters: Vec<CrossDocCluster>,
296}
297
298impl Topic {
299    /// Create a new topic.
300    pub fn new(id: &str) -> Self {
301        Self {
302            id: id.to_string(),
303            documents: Vec::new(),
304            gold_clusters: Vec::new(),
305        }
306    }
307
308    /// Add a document.
309    pub fn add_document(&mut self, doc: Document) {
310        self.documents.push(doc);
311    }
312
313    /// Add a gold cluster.
314    pub fn add_gold_cluster(&mut self, cluster: CrossDocCluster) {
315        self.gold_clusters.push(cluster);
316    }
317}
318
319/// A long document with gold annotations (for long-document evaluation).
320#[derive(Debug, Clone)]
321pub struct LongDocument {
322    /// Document ID
323    pub id: String,
324    /// Full text
325    pub text: String,
326    /// Gold coreference chains
327    pub gold_chains: Vec<CorefChain>,
328    /// Optional: pre-computed windows
329    pub windows: Option<Vec<WindowOutput>>,
330}
331
332impl LongDocument {
333    /// Create a new long document.
334    pub fn new(id: &str, text: &str, gold_chains: Vec<CorefChain>) -> Self {
335        Self {
336            id: id.to_string(),
337            text: text.to_string(),
338            gold_chains,
339            windows: None,
340        }
341    }
342
343    /// Get text length in characters.
344    pub fn char_len(&self) -> usize {
345        self.text.chars().count()
346    }
347
348    /// Estimate token count (rough: chars / 5).
349    pub fn approx_tokens(&self) -> usize {
350        self.text.len() / 5
351    }
352}
353
354// =============================================================================
355// Evaluation Functions
356// =============================================================================
357
358/// Evaluate cross-document coreference on a set of topics.
359///
360/// Uses the `UnifiedCrossContextResolver` to merge clusters across documents.
361pub fn evaluate_cross_document<E: ClusterEncoder + Clone, S: MergeScorer + Clone>(
362    topics: &[Topic],
363    encoder: E,
364    scorer: S,
365    config: &CrossContextEvalConfig,
366) -> Result<CrossContextEvalResults> {
367    let start = std::time::Instant::now();
368
369    let resolver_config = CrossContextConfig {
370        window_size: config.window_size,
371        window_overlap: config.window_overlap,
372        merge_threshold: config.merge_threshold,
373    };
374
375    let resolver = UnifiedCrossContextResolver::new(encoder, scorer, resolver_config);
376
377    let mut all_gold_chains: Vec<CorefChain> = Vec::new();
378    let mut all_pred_chains: Vec<CorefChain> = Vec::new();
379    let mut per_topic = HashMap::new();
380    let mut total_gold_clusters = 0;
381    let mut total_pred_clusters = 0;
382
383    for topic in topics {
384        // Convert gold clusters to chains for evaluation
385        let topic_gold_chains =
386            cross_doc_clusters_to_chains(&topic.gold_clusters, &topic.documents);
387        total_gold_clusters += topic.gold_clusters.len();
388
389        // Resolve across documents in this topic
390        let pred_clusters = resolver.resolve_documents(&topic.documents);
391        total_pred_clusters += pred_clusters.len();
392
393        let topic_pred_chains = cross_doc_clusters_to_chains(&pred_clusters, &topic.documents);
394
395        // Compute per-topic metrics
396        let topic_f1 = conll_f1(&topic_gold_chains, &topic_pred_chains);
397
398        per_topic.insert(
399            topic.id.clone(),
400            TopicResults {
401                topic_id: topic.id.clone(),
402                num_documents: topic.documents.len(),
403                conll_f1: topic_f1,
404                num_gold_clusters: topic.gold_clusters.len(),
405                num_pred_clusters: pred_clusters.len(),
406            },
407        );
408
409        all_gold_chains.extend(topic_gold_chains);
410        all_pred_chains.extend(topic_pred_chains);
411    }
412
413    // Compute aggregate metrics
414    let (muc_p, muc_r, muc_f1) =
415        crate::eval::coref_metrics::muc_score(&all_pred_chains, &all_gold_chains);
416    let (b3_p, b3_r, b3_f1) =
417        crate::eval::coref_metrics::b_cubed_score(&all_pred_chains, &all_gold_chains);
418    let (ceaf_p, ceaf_r, ceaf_f1) =
419        crate::eval::coref_metrics::ceaf_e_score(&all_pred_chains, &all_gold_chains);
420    let (lea_p, lea_r, lea_f1) =
421        crate::eval::coref_metrics::lea_score(&all_pred_chains, &all_gold_chains);
422    let conll = conll_f1(&all_gold_chains, &all_pred_chains);
423
424    let num_contexts: usize = topics.iter().map(|t| t.documents.len()).sum();
425    let total_mentions: usize = all_pred_chains.iter().map(|c| c.len()).sum();
426    let avg_cluster_size = if !all_pred_chains.is_empty() {
427        total_mentions as f64 / all_pred_chains.len() as f64
428    } else {
429        0.0
430    };
431
432    Ok(CrossContextEvalResults {
433        benchmark: "Cross-Document".to_string(),
434        config: config.clone(),
435        muc: CorefScores::from_tuple((muc_p, muc_r, muc_f1)),
436        b_cubed: CorefScores::from_tuple((b3_p, b3_r, b3_f1)),
437        ceaf_e: CorefScores::from_tuple((ceaf_p, ceaf_r, ceaf_f1)),
438        lea: CorefScores::from_tuple((lea_p, lea_r, lea_f1)),
439        conll_f1: conll,
440        num_contexts,
441        num_gold_clusters: total_gold_clusters,
442        num_pred_clusters: total_pred_clusters,
443        avg_cluster_size,
444        time_ms: start.elapsed().as_millis() as f64,
445        per_topic: Some(per_topic),
446        per_document: None,
447    })
448}
449
450/// Evaluate long-document coreference.
451///
452/// Uses the `UnifiedCrossContextResolver` to merge clusters across windows.
453pub fn evaluate_long_document<E: ClusterEncoder + Clone, S: MergeScorer + Clone>(
454    documents: &[LongDocument],
455    encoder: E,
456    scorer: S,
457    config: &CrossContextEvalConfig,
458) -> Result<CrossContextEvalResults> {
459    let start = std::time::Instant::now();
460
461    let resolver_config = CrossContextConfig {
462        window_size: config.window_size,
463        window_overlap: config.window_overlap,
464        merge_threshold: config.merge_threshold,
465    };
466
467    let resolver = UnifiedCrossContextResolver::new(encoder, scorer, resolver_config);
468
469    let mut all_gold_chains: Vec<CorefChain> = Vec::new();
470    let mut all_pred_chains: Vec<CorefChain> = Vec::new();
471    let mut per_document = HashMap::new();
472
473    for doc in documents {
474        // Use pre-computed windows if available, otherwise would need to compute
475        let windows = doc.windows.clone().unwrap_or_default();
476
477        if windows.is_empty() {
478            // No pre-computed windows: treat the entire document as a single window
479            // and let the resolver merge from that single context.
480            let single_window = WindowOutput::new(
481                0,
482                0,
483                doc.char_len(),
484                if config.use_gold_mentions {
485                    doc.gold_chains.clone()
486                } else {
487                    // Without gold mentions we have no mention detector here;
488                    // produce an empty prediction so metrics reflect the gap.
489                    Vec::new()
490                },
491            );
492            let pred_chains = resolver.resolve_long_document_windows(&[single_window]);
493
494            let doc_f1 = conll_f1(&doc.gold_chains, &pred_chains);
495            per_document.insert(
496                doc.id.clone(),
497                DocumentResults {
498                    doc_id: doc.id.clone(),
499                    num_tokens: doc.approx_tokens(),
500                    num_windows: 1,
501                    conll_f1: doc_f1,
502                    num_gold_chains: doc.gold_chains.len(),
503                    num_pred_chains: pred_chains.len(),
504                },
505            );
506
507            all_gold_chains.extend(doc.gold_chains.clone());
508            all_pred_chains.extend(pred_chains);
509            continue;
510        }
511
512        let pred_chains = resolver.resolve_long_document_windows(&windows);
513
514        // Compute per-document metrics
515        let doc_f1 = conll_f1(&doc.gold_chains, &pred_chains);
516
517        per_document.insert(
518            doc.id.clone(),
519            DocumentResults {
520                doc_id: doc.id.clone(),
521                num_tokens: doc.approx_tokens(),
522                num_windows: windows.len(),
523                conll_f1: doc_f1,
524                num_gold_chains: doc.gold_chains.len(),
525                num_pred_chains: pred_chains.len(),
526            },
527        );
528
529        all_gold_chains.extend(doc.gold_chains.clone());
530        all_pred_chains.extend(pred_chains);
531    }
532
533    // Compute aggregate metrics
534    let (muc_p, muc_r, muc_f1) =
535        crate::eval::coref_metrics::muc_score(&all_pred_chains, &all_gold_chains);
536    let (b3_p, b3_r, b3_f1) =
537        crate::eval::coref_metrics::b_cubed_score(&all_pred_chains, &all_gold_chains);
538    let (ceaf_p, ceaf_r, ceaf_f1) =
539        crate::eval::coref_metrics::ceaf_e_score(&all_pred_chains, &all_gold_chains);
540    let (lea_p, lea_r, lea_f1) =
541        crate::eval::coref_metrics::lea_score(&all_pred_chains, &all_gold_chains);
542    let conll = conll_f1(&all_gold_chains, &all_pred_chains);
543
544    let total_mentions: usize = all_pred_chains.iter().map(|c| c.len()).sum();
545    let avg_cluster_size = if !all_pred_chains.is_empty() {
546        total_mentions as f64 / all_pred_chains.len() as f64
547    } else {
548        0.0
549    };
550
551    Ok(CrossContextEvalResults {
552        benchmark: "Long-Document".to_string(),
553        config: config.clone(),
554        muc: CorefScores::from_tuple((muc_p, muc_r, muc_f1)),
555        b_cubed: CorefScores::from_tuple((b3_p, b3_r, b3_f1)),
556        ceaf_e: CorefScores::from_tuple((ceaf_p, ceaf_r, ceaf_f1)),
557        lea: CorefScores::from_tuple((lea_p, lea_r, lea_f1)),
558        conll_f1: conll,
559        num_contexts: documents.len(),
560        num_gold_clusters: all_gold_chains.len(),
561        num_pred_clusters: all_pred_chains.len(),
562        avg_cluster_size,
563        time_ms: start.elapsed().as_millis() as f64,
564        per_topic: None,
565        per_document: Some(per_document),
566    })
567}
568
569// =============================================================================
570// Helper Functions
571// =============================================================================
572
573/// Convert cross-document clusters to coreference chains.
574fn cross_doc_clusters_to_chains(
575    clusters: &[CrossDocCluster],
576    docs: &[Document],
577) -> Vec<CorefChain> {
578    clusters
579        .iter()
580        .map(|cluster| {
581            let mentions: Vec<Mention> = cluster
582                .mentions
583                .iter()
584                .filter_map(|(doc_id, entity_idx)| {
585                    let doc = docs.iter().find(|d| &d.id == doc_id)?;
586                    let entity = doc.entities.get(*entity_idx)?;
587                    Some(Mention {
588                        text: entity.text.clone(),
589                        start: entity.start(),
590                        end: entity.end(),
591                        head_start: None,
592                        head_end: None,
593                        entity_type: Some(entity.entity_type.as_label().to_string()),
594                        mention_type: None,
595                    })
596                })
597                .collect();
598            CorefChain::new(mentions)
599        })
600        .filter(|c| !c.is_empty())
601        .collect()
602}
603
604// =============================================================================
605// Stepwise Error Analysis (Table 5 from xCoRe paper)
606// =============================================================================
607
608/// Stepwise error analysis configuration.
609#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
610pub enum StepwiseAnalysis {
611    /// Full pipeline (predicted mentions, predicted clusters)
612    FullPipeline,
613    /// Gold mentions, predicted clusters
614    GoldMentions,
615    /// Gold mentions, gold clusters (cluster merging only)
616    GoldMentionsAndClusters,
617}
618
619impl StepwiseAnalysis {
620    /// Get description.
621    pub fn description(&self) -> &'static str {
622        match self {
623            Self::FullPipeline => "xCoRe (full pipeline)",
624            Self::GoldMentions => "xCoRe (gold mentions)",
625            Self::GoldMentionsAndClusters => "xCoRe (gold mentions & clusters)",
626        }
627    }
628}
629
630/// Run stepwise error analysis as in xCoRe Table 5.
631///
632/// This helps identify which pipeline stage is the bottleneck:
633/// - Mention extraction
634/// - Within-context clustering
635/// - Cross-context cluster merging
636pub fn stepwise_error_analysis<E: ClusterEncoder + Clone, S: MergeScorer + Clone>(
637    benchmark: CrossContextBenchmark,
638    topics: &[Topic],           // For cross-doc
639    documents: &[LongDocument], // For long-doc
640    encoder: E,
641    scorer: S,
642) -> Result<HashMap<StepwiseAnalysis, CrossContextEvalResults>> {
643    let mut results = HashMap::new();
644
645    for analysis in [
646        StepwiseAnalysis::FullPipeline,
647        StepwiseAnalysis::GoldMentions,
648        StepwiseAnalysis::GoldMentionsAndClusters,
649    ] {
650        let config = match analysis {
651            StepwiseAnalysis::FullPipeline => CrossContextEvalConfig::for_benchmark(benchmark),
652            StepwiseAnalysis::GoldMentions => CrossContextEvalConfig::gold_clusters(),
653            StepwiseAnalysis::GoldMentionsAndClusters => CrossContextEvalConfig::oracle(),
654        };
655
656        let eval_result = if benchmark.is_cross_document() {
657            evaluate_cross_document(topics, encoder.clone(), scorer.clone(), &config)?
658        } else {
659            evaluate_long_document(documents, encoder.clone(), scorer.clone(), &config)?
660        };
661
662        results.insert(analysis, eval_result);
663    }
664
665    Ok(results)
666}
667
668// =============================================================================
669// Tests
670// =============================================================================
671
672#[cfg(test)]
673mod tests {
674    use super::*;
675    use crate::eval::cluster_encoder::{CosineMergeScorer, HeuristicClusterEncoder};
676    use anno::{Entity, EntityType};
677
678    #[test]
679    fn test_benchmark_properties() {
680        assert!(CrossContextBenchmark::ECBPlus.is_cross_document());
681        assert!(!CrossContextBenchmark::ECBPlus.is_long_document());
682
683        assert!(!CrossContextBenchmark::LitBank.is_cross_document());
684        assert!(CrossContextBenchmark::LitBank.is_long_document());
685
686        assert_eq!(CrossContextBenchmark::all().len(), 5);
687        assert_eq!(CrossContextBenchmark::cross_document().len(), 2);
688        assert_eq!(CrossContextBenchmark::long_document().len(), 3);
689    }
690
691    #[test]
692    fn test_benchmark_sota() {
693        assert!((CrossContextBenchmark::ECBPlus.xcore_sota_f1() - 40.3).abs() < 0.1);
694        assert!((CrossContextBenchmark::LitBank.xcore_sota_f1() - 78.2).abs() < 0.1);
695    }
696
697    #[test]
698    fn test_eval_config_default() {
699        let config = CrossContextEvalConfig::default();
700        assert_eq!(config.window_size, 4000);
701        assert!(!config.use_gold_mentions);
702    }
703
704    #[test]
705    fn test_eval_config_for_benchmark() {
706        let config = CrossContextEvalConfig::for_benchmark(CrossContextBenchmark::ECBPlus);
707        assert_eq!(config.window_size, 512);
708    }
709
710    #[test]
711    fn test_topic_creation() {
712        let mut topic = Topic::new("topic_1");
713        topic.add_document(Document::new("doc1", "Obama visited Paris."));
714        topic.add_document(Document::new("doc2", "The president met leaders."));
715
716        assert_eq!(topic.documents.len(), 2);
717    }
718
719    #[test]
720    fn test_long_document_creation() {
721        use anno::MentionType;
722
723        fn new_mention(text: &str, start: usize, end: usize) -> Mention {
724            Mention {
725                text: text.to_string(),
726                start,
727                end,
728                head_start: None,
729                head_end: None,
730                entity_type: None,
731                mention_type: Some(MentionType::Proper),
732            }
733        }
734
735        let chains = vec![CorefChain::new(vec![
736            new_mention("Obama", 0, 5),
737            new_mention("he", 100, 102),
738        ])];
739
740        let doc = LongDocument::new(
741            "book1",
742            "Obama went to Paris. ".repeat(100).as_str(),
743            chains,
744        );
745
746        assert!(doc.approx_tokens() > 100);
747        assert_eq!(doc.gold_chains.len(), 1);
748    }
749
750    #[test]
751    fn test_evaluate_cross_document_empty() {
752        let encoder = HeuristicClusterEncoder::new(64);
753        let scorer = CosineMergeScorer::new();
754        let config = CrossContextEvalConfig::default();
755
756        let topics: Vec<Topic> = vec![];
757        let results = evaluate_cross_document(&topics, encoder, scorer, &config).unwrap();
758
759        assert_eq!(results.num_contexts, 0);
760    }
761
762    #[test]
763    fn test_evaluate_cross_document_single_topic() {
764        let encoder = HeuristicClusterEncoder::new(64);
765        let scorer = CosineMergeScorer::new();
766        let config = CrossContextEvalConfig::default();
767
768        let mut topic = Topic::new("topic_1");
769        topic.add_document(
770            Document::new("doc1", "Obama visited France.").with_entities(vec![Entity::new(
771                "Obama",
772                EntityType::Person,
773                0,
774                5,
775                0.9,
776            )]),
777        );
778        topic.add_document(
779            Document::new("doc2", "The president met Macron.").with_entities(vec![
780                Entity::new("The president", EntityType::Person, 0, 13, 0.8),
781                Entity::new("Macron", EntityType::Person, 18, 24, 0.9),
782            ]),
783        );
784
785        let results = evaluate_cross_document(&[topic], encoder, scorer, &config).unwrap();
786
787        assert_eq!(results.num_contexts, 2);
788        assert!(results.per_topic.is_some());
789    }
790
791    #[test]
792    fn test_stepwise_analysis_types() {
793        assert_eq!(
794            StepwiseAnalysis::FullPipeline.description(),
795            "xCoRe (full pipeline)"
796        );
797        assert_eq!(
798            StepwiseAnalysis::GoldMentions.description(),
799            "xCoRe (gold mentions)"
800        );
801    }
802
803    #[test]
804    fn test_results_summary() {
805        let results = CrossContextEvalResults {
806            benchmark: "Test".to_string(),
807            config: CrossContextEvalConfig::default(),
808            muc: CorefScores::new(0.8, 0.7),
809            b_cubed: CorefScores::new(0.75, 0.65),
810            ceaf_e: CorefScores::new(0.7, 0.6),
811            lea: CorefScores::new(0.72, 0.62),
812            conll_f1: 0.70,
813            num_contexts: 10,
814            num_gold_clusters: 50,
815            num_pred_clusters: 45,
816            avg_cluster_size: 2.5,
817            time_ms: 100.0,
818            per_topic: None,
819            per_document: None,
820        };
821
822        let summary = results.summary();
823        assert!(summary.contains("70.0%"));
824    }
825
826    #[test]
827    fn test_evaluate_cross_document_with_synthetic_data() {
828        // 2 topics with 3 docs each, overlapping entity names across docs
829        let encoder = HeuristicClusterEncoder::new(64);
830        let scorer = CosineMergeScorer::new();
831        let config = CrossContextEvalConfig::default();
832
833        // Topic 1: Obama visits France
834        let mut topic1 = Topic::new("politics");
835        topic1.add_document(
836            Document::new("doc1", "Obama visited France yesterday.").with_entities(vec![
837                Entity::new("Obama", EntityType::Person, 0, 5, 0.9),
838                Entity::new("France", EntityType::Location, 14, 20, 0.9),
839            ]),
840        );
841        topic1.add_document(
842            Document::new("doc2", "The president arrived in Paris.").with_entities(vec![
843                Entity::new("The president", EntityType::Person, 0, 13, 0.8),
844                Entity::new("Paris", EntityType::Location, 25, 30, 0.9),
845            ]),
846        );
847        topic1.add_document(
848            Document::new("doc3", "Barack Obama met Macron in France.").with_entities(vec![
849                Entity::new("Barack Obama", EntityType::Person, 0, 12, 0.95),
850                Entity::new("Macron", EntityType::Person, 17, 23, 0.9),
851                Entity::new("France", EntityType::Location, 27, 33, 0.9),
852            ]),
853        );
854        // Gold: Obama cluster across 3 docs, France/Paris cluster across 2 docs
855        let mut obama_cluster = crate::eval::cdcr::CrossDocCluster::new(0u64, "Obama");
856        obama_cluster.mentions = vec![
857            ("doc1".to_string(), 0),
858            ("doc2".to_string(), 0),
859            ("doc3".to_string(), 0),
860        ];
861        let mut france_cluster = crate::eval::cdcr::CrossDocCluster::new(1u64, "France");
862        france_cluster.mentions = vec![
863            ("doc1".to_string(), 1),
864            ("doc2".to_string(), 1),
865            ("doc3".to_string(), 2),
866        ];
867        topic1.add_gold_cluster(obama_cluster);
868        topic1.add_gold_cluster(france_cluster);
869
870        // Topic 2: Tech companies
871        let mut topic2 = Topic::new("tech");
872        topic2.add_document(
873            Document::new("doc4", "Apple released new products.").with_entities(vec![Entity::new(
874                "Apple",
875                EntityType::Organization,
876                0,
877                5,
878                0.9,
879            )]),
880        );
881        topic2.add_document(
882            Document::new("doc5", "The company expanded in Asia.").with_entities(vec![
883                Entity::new("The company", EntityType::Organization, 0, 11, 0.8),
884                Entity::new("Asia", EntityType::Location, 24, 28, 0.9),
885            ]),
886        );
887        topic2.add_document(
888            Document::new("doc6", "Apple Inc announced quarterly results.").with_entities(vec![
889                Entity::new("Apple Inc", EntityType::Organization, 0, 9, 0.9),
890            ]),
891        );
892        let mut apple_cluster = crate::eval::cdcr::CrossDocCluster::new(0u64, "Apple");
893        apple_cluster.mentions = vec![
894            ("doc4".to_string(), 0),
895            ("doc5".to_string(), 0),
896            ("doc6".to_string(), 0),
897        ];
898        topic2.add_gold_cluster(apple_cluster);
899
900        let results = evaluate_cross_document(&[topic1, topic2], encoder, scorer, &config).unwrap();
901
902        // Basic sanity: metrics in valid ranges
903        assert!(results.conll_f1 >= 0.0 && results.conll_f1 <= 1.0);
904        assert!(results.muc.f1 >= 0.0 && results.muc.f1 <= 1.0);
905        assert!(results.b_cubed.f1 >= 0.0 && results.b_cubed.f1 <= 1.0);
906        assert!(results.ceaf_e.f1 >= 0.0 && results.ceaf_e.f1 <= 1.0);
907
908        // Should have evaluated 6 documents across 2 topics
909        assert_eq!(results.num_contexts, 6);
910        assert!(results.per_topic.is_some());
911        let per_topic = results.per_topic.as_ref().unwrap();
912        assert_eq!(per_topic.len(), 2);
913        assert!(per_topic.contains_key("politics"));
914        assert!(per_topic.contains_key("tech"));
915
916        // Gold clusters: 3 total (obama, france, apple)
917        assert_eq!(results.num_gold_clusters, 3);
918
919        // Predicted clusters should be non-zero (heuristic encoder finds something)
920        assert!(results.num_pred_clusters > 0);
921    }
922
923    #[test]
924    fn test_evaluate_long_document_with_gold_mentions() {
925        let encoder = HeuristicClusterEncoder::new(64);
926        let scorer = CosineMergeScorer::new();
927        let config = CrossContextEvalConfig {
928            use_gold_mentions: true,
929            ..CrossContextEvalConfig::default()
930        };
931
932        let chains = vec![
933            CorefChain::new(vec![
934                Mention::new("Obama", 0, 5),
935                Mention::new("he", 50, 52),
936            ]),
937            CorefChain::new(vec![
938                Mention::new("France", 14, 20),
939                Mention::new("the country", 60, 71),
940            ]),
941        ];
942
943        let doc = LongDocument::new("long_doc", &"Obama visited France. ".repeat(10), chains);
944
945        let results = evaluate_long_document(&[doc], encoder, scorer, &config).unwrap();
946        assert_eq!(results.num_contexts, 1);
947        // With gold mentions in single-window mode, should produce per-doc results
948        assert!(results.per_document.is_some());
949        let per_doc = results.per_document.as_ref().unwrap();
950        assert_eq!(per_doc.len(), 1);
951        assert!(per_doc.contains_key("long_doc"));
952    }
953}
anno_eval/eval/cross_context_eval.rs

anno_eval/eval/
cross_context_eval.rs
