fuzzy_align.py (4596B)
1 import re 2 import numpy 3 4 5 BREAK_PENALTY = 10 # Penalty if variable is not preceded by space or followed by space/period. 6 TEMPLATE_MARKER = '\1' # Template marker for variables 7 POST_TEMPLATE_VALID = " .," # Unpenalized characters after template variable 8 PRE_TEMPLATE_VALID = " " # Unpenalized characters before template variable 9 10 11 class FuzzyAlign: 12 """ 13 Use a Levenshtein algorithm to fuzzily align the output of a model with its template. 14 15 As currently written, the variables must always be preceded by a space in the template. 16 Furthermore, the following character is expected to be one of POST_TEMPLATE_VALID. 17 In any case, when matching, a PRE_TEMPLATE_VALID will be enforced before a variable and any POST_TEMPLATE_VALID character will match the template just after a variable. 18 """ 19 def __init__(self, template): 20 self.fields = re.findall(r"{([^}]+)}", template) 21 self.template = re.sub(" {[^}]*}", TEMPLATE_MARKER, template) 22 23 def __call__(self, prediction): 24 m = numpy.empty((len(prediction)+1, len(self.template)+1), dtype=numpy.int64) 25 parents = numpy.zeros((len(prediction)+1, len(self.template)+1), dtype=numpy.int8) 26 parents[0, 1:] |= 2 27 parents[1:, 0] |= 1 28 29 m[:, 0] = numpy.arange(len(prediction)+1) 30 31 template_init = (numpy.array(list(self.template))==TEMPLATE_MARKER)*(BREAK_PENALTY-1)+1 32 m[0, 1:] = template_init.cumsum() 33 34 35 for y in range(1, len(prediction)+1): 36 for x in range(1, len(self.template)+1): 37 if x>1 and self.template[x-2] == TEMPLATE_MARKER: 38 addition = 1 39 deletion = BREAK_PENALTY 40 transformation = (0 if prediction[y-1] in POST_TEMPLATE_VALID else BREAK_PENALTY) 41 elif self.template[x-1] == TEMPLATE_MARKER: 42 addition = 0 43 deletion = BREAK_PENALTY 44 transformation = (0 if prediction[y-1] in PRE_TEMPLATE_VALID else BREAK_PENALTY) 45 else: 46 addition = 1 47 deletion = 1 48 transformation = (self.template[x-1] != prediction[y-1]) 49 50 m[y][x] = min(m[y-1][x] + addition, m[y][x-1] + deletion, m[y-1][x-1] + transformation) 51 if m[y-1][x] + addition == m[y][x]: 52 parents[y][x] |= 1 53 if m[y][x-1] + deletion == m[y][x]: 54 parents[y][x] |= 2 55 if m[y-1][x-1] + transformation == m[y][x]: 56 parents[y][x] |= 4 57 58 matches = {location+1: [] for location, value in enumerate(self.template) if value == TEMPLATE_MARKER} 59 y, x = len(prediction), len(self.template) 60 while (y, x) != (0, 0): 61 if self.template[x-1] == TEMPLATE_MARKER: 62 if parents[y][x] & 1: 63 y -= 1 64 matches[x].append(prediction[y]) 65 elif parents[y][x] & 4: 66 y -= 1 67 x -= 1 68 elif parents[y][x] & 2: 69 x -= 1 70 else: 71 if parents[y][x] & 4: 72 y -= 1 73 x -= 1 74 elif parents[y][x] & 1: 75 y -= 1 76 elif parents[y][x] & 2: 77 x -= 1 78 79 return { field: "".join(match[::-1]).strip() for field, (_, match) in zip(self.fields, sorted(matches.items())) } 80 81 82 if __name__ == "__main__": 83 prediction = "The conflict between Government of Syria (A) and IS (B) resulted in causalities between 2020-02-22 and 2020-01-22 at Al Baghouz town in Albu Kamal district in the region of Deir ez Zor governorate in Syria. 0 people died on side A, 1 people died in side B, 0 civilians died and 0 unidentified people died. The low estimate for the total number of people that died is 1. The high estimate for this total number is 1 people." 84 template = "The conflict between {side_a_name} (A) and {side_b_name} (B) resulted in causalities between {start_date} and {end_date} at {location_where_name} in {location_adm2_name} in the region of {location_adm1_name} in {location_root_name}. {deaths_side_a} people died on side A, {deaths_side_b} people died on side B, {deaths_civilian} civilians died and {deaths_unknown} unidentified people died. The low estimate for the total number of people that died is {deaths_low}. The high estimate for the total number of people that died is {deaths_high}." 85 86 print("\n".join(f"{key}: |{value}|" for key, value in FuzzyAlign(template)(prediction).items()))