@@ -16,8 +16,9 @@ class ConllCorefScores(Metric):
1616
1717 supports_distributed = True
1818
19- def __init__ (self ) -> None :
19+ def __init__ (self , allow_singletons = False ) -> None :
2020 self .scorers = [Scorer (m ) for m in (Scorer .muc , Scorer .b_cubed , Scorer .ceafe )]
21+ self .allow_singletons = allow_singletons
2122
2223 @overrides
2324 def __call__ (
@@ -56,7 +57,7 @@ def __call__(
5657 for i , metadata in enumerate (metadata_list ):
5758 gold_clusters , mention_to_gold = self .get_gold_clusters (metadata ["clusters" ])
5859 predicted_clusters , mention_to_predicted = self .get_predicted_clusters (
59- top_spans [i ], antecedent_indices [i ], predicted_antecedents [i ]
60+ top_spans [i ], antecedent_indices [i ], predicted_antecedents [i ], self . allow_singletons
6061 )
6162 for scorer in self .scorers :
6263 scorer .update (
@@ -91,6 +92,7 @@ def get_predicted_clusters(
9192 top_spans : torch .Tensor , # (num_spans, 2)
9293 antecedent_indices : torch .Tensor , # (num_spans, num_antecedents)
9394 predicted_antecedents : torch .Tensor , # (num_spans,)
95+ allow_singletons : bool ,
9496 ) -> Tuple [
9597 List [Tuple [Tuple [int , int ], ...]], Dict [Tuple [int , int ], Tuple [Tuple [int , int ], ...]]
9698 ]:
@@ -104,7 +106,10 @@ def get_predicted_clusters(
104106 # Find predicted index in the antecedent spans.
105107 predicted_index = antecedent_indices [i , predicted_antecedent ]
106108 # Must be a previous span.
107- assert i > predicted_index
109+ if allow_singletons :
110+ assert i >= predicted_index
111+ else :
112+ assert i > predicted_index
108113 antecedent_span : Tuple [int , int ] = tuple ( # type: ignore
109114 top_spans [predicted_index ].tolist ()
110115 )
0 commit comments