1use anno::{DiscontinuousEntity, DiscontinuousNER, RelationExtractor, Result};
27use serde::{Deserialize, Serialize};
28use std::collections::HashMap;
29
30use super::dataset::synthetic::{
31 discontinuous::dataset as discontinuous_dataset, relations::dataset as relations_dataset,
32};
33use super::discontinuous::{
34 evaluate_discontinuous_ner, DiscontinuousEvalConfig, DiscontinuousGold, DiscontinuousNERMetrics,
35};
36use super::relation::{
37 evaluate_relations, RelationEvalConfig, RelationGold, RelationMetrics, RelationPrediction,
38};
39use super::visual::{
40 evaluate_visual_ner, synthetic_visual_examples, VisualEvalConfig, VisualGold, VisualNERMetrics,
41 VisualPrediction,
42};
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct AdvancedTaskResults {
51 pub timestamp: String,
53 pub task: String,
55 pub models: Vec<ModelResult>,
57 pub num_examples: usize,
59 pub num_gold: usize,
61}
62
63impl AdvancedTaskResults {
64 pub fn summary(&self) -> String {
66 let mut s = format!(
67 "=== {} Evaluation ({} examples) ===\n",
68 self.task, self.num_examples
69 );
70
71 for model in &self.models {
72 s.push_str(&format!(
73 "\n{}: F1={:.1}%\n",
74 model.name,
75 model.primary_f1 * 100.0
76 ));
77 }
78
79 s
80 }
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct ModelResult {
86 pub name: String,
88 pub primary_f1: f64,
90 pub metrics: HashMap<String, f64>,
92}
93
94pub fn evaluate_discontinuous_synthetic<M: DiscontinuousNER>(
102 model: &M,
103 labels: &[&str],
104 threshold: f32,
105) -> Result<DiscontinuousNERMetrics> {
106 let examples = discontinuous_dataset();
107 let config = DiscontinuousEvalConfig::default();
108
109 let mut all_gold: Vec<DiscontinuousGold> = Vec::new();
110 let mut all_pred: Vec<DiscontinuousEntity> = Vec::new();
111
112 for example in &examples {
113 all_gold.extend(example.entities.clone());
115
116 let pred = model.extract_discontinuous(&example.text, labels, threshold)?;
118 all_pred.extend(pred);
119 }
120
121 Ok(evaluate_discontinuous_ner(&all_gold, &all_pred, &config))
122}
123
124pub fn evaluate_discontinuous_gold_vs_gold() -> DiscontinuousNERMetrics {
126 let examples = discontinuous_dataset();
127 let config = DiscontinuousEvalConfig::default();
128
129 let gold: Vec<DiscontinuousGold> = examples.iter().flat_map(|ex| ex.entities.clone()).collect();
130
131 let pred: Vec<DiscontinuousEntity> = gold
133 .iter()
134 .map(|g| DiscontinuousEntity {
135 spans: g.spans.clone(),
136 text: g.text.clone(),
137 entity_type: g.entity_type.clone(),
138 confidence: anno::Confidence::ONE,
139 })
140 .collect();
141
142 evaluate_discontinuous_ner(&gold, &pred, &config)
143}
144
145pub fn evaluate_relations_synthetic<M: RelationExtractor>(
151 model: &M,
152 labels: &[&str],
153 relations: &[&str],
154 threshold: f32,
155) -> Result<RelationMetrics> {
156 let examples = relations_dataset();
157 let config = RelationEvalConfig::default();
158
159 let mut all_gold: Vec<RelationGold> = Vec::new();
160 let mut all_pred: Vec<RelationPrediction> = Vec::new();
161
162 for example in &examples {
163 all_gold.extend(example.relations.clone());
165
166 let result = model.extract_with_relations(&example.text, labels, relations, threshold)?;
168
169 for rel in &result.relations {
171 if rel.head_idx < result.entities.len() && rel.tail_idx < result.entities.len() {
172 let head = &result.entities[rel.head_idx];
173 let tail = &result.entities[rel.tail_idx];
174 all_pred.push(RelationPrediction {
175 head_span: (head.start(), head.end()),
176 head_type: head.entity_type.as_label().to_string(),
177 tail_span: (tail.start(), tail.end()),
178 tail_type: tail.entity_type.as_label().to_string(),
179 relation_type: rel.relation_type.clone(),
180 confidence: rel.confidence.value() as f32,
181 });
182 }
183 }
184 }
185
186 Ok(evaluate_relations(&all_gold, &all_pred, &config))
187}
188
189pub fn evaluate_relations_gold_vs_gold() -> RelationMetrics {
191 let examples = relations_dataset();
192 let config = RelationEvalConfig::default();
193
194 let gold: Vec<RelationGold> = examples
195 .iter()
196 .flat_map(|ex| ex.relations.clone())
197 .collect();
198
199 let pred: Vec<RelationPrediction> = gold
201 .iter()
202 .map(|g| RelationPrediction {
203 head_span: g.head_span,
204 head_type: g.head_type.clone(),
205 tail_span: g.tail_span,
206 tail_type: g.tail_type.clone(),
207 relation_type: g.relation_type.clone(),
208 confidence: 1.0,
209 })
210 .collect();
211
212 evaluate_relations(&gold, &pred, &config)
213}
214
215pub fn evaluate_visual_gold_vs_gold() -> VisualNERMetrics {
221 let examples = synthetic_visual_examples();
222 let config = VisualEvalConfig::default();
223
224 let gold: Vec<VisualGold> = examples
225 .iter()
226 .flat_map(|(_, entities)| entities.clone())
227 .collect();
228
229 let pred: Vec<VisualPrediction> = gold
231 .iter()
232 .map(|g| VisualPrediction {
233 text: g.text.clone(),
234 entity_type: g.entity_type.clone(),
235 bbox: g.bbox,
236 confidence: 1.0,
237 })
238 .collect();
239
240 evaluate_visual_ner(&gold, &pred, &config)
241}
242
243pub fn synthetic_dataset_stats() -> SyntheticDatasetStats {
249 let disc = discontinuous_dataset();
250 let rel = relations_dataset();
251 let vis = synthetic_visual_examples();
252
253 SyntheticDatasetStats {
254 discontinuous_examples: disc.len(),
255 discontinuous_entities: disc.iter().map(|ex| ex.entities.len()).sum(),
256 relation_examples: rel.len(),
257 relations: rel.iter().map(|ex| ex.relations.len()).sum(),
258 visual_examples: vis.len(),
259 visual_entities: vis.iter().map(|(_, e)| e.len()).sum(),
260 }
261}
262
263#[derive(Debug, Clone, Serialize, Deserialize)]
265pub struct SyntheticDatasetStats {
266 pub discontinuous_examples: usize,
268 pub discontinuous_entities: usize,
270 pub relation_examples: usize,
272 pub relations: usize,
274 pub visual_examples: usize,
276 pub visual_entities: usize,
278}
279
280#[cfg(test)]
285mod tests {
286 use super::*;
287
288 #[test]
289 fn test_discontinuous_gold_vs_gold() {
290 let metrics = evaluate_discontinuous_gold_vs_gold();
291 assert!(
292 (metrics.exact_f1 - 1.0).abs() < 0.001,
293 "Perfect prediction should give F1=1.0, got {}",
294 metrics.exact_f1
295 );
296 }
297
298 #[test]
299 fn test_relations_gold_vs_gold() {
300 let metrics = evaluate_relations_gold_vs_gold();
301 assert!(
302 (metrics.strict_f1 - 1.0).abs() < 0.001,
303 "Perfect prediction should give F1=1.0, got {}",
304 metrics.strict_f1
305 );
306 }
307
308 #[test]
309 fn test_visual_gold_vs_gold() {
310 let metrics = evaluate_visual_gold_vs_gold();
311 assert!(
312 (metrics.e2e_f1 - 1.0).abs() < 0.001,
313 "Perfect prediction should give F1=1.0, got {}",
314 metrics.e2e_f1
315 );
316 }
317
318 #[test]
319 fn test_synthetic_dataset_stats() {
320 let stats = synthetic_dataset_stats();
321 assert!(stats.discontinuous_examples > 0);
322 assert!(stats.discontinuous_entities > 0);
323 assert!(stats.relation_examples > 0);
324 assert!(stats.relations > 0);
325 assert!(stats.visual_examples > 0);
326 assert!(stats.visual_entities > 0);
327 }
328}