ml6team/post-processing-summarization

1from typing import Dict2from PIL import ImageFont3 4TPL_DEP_WORDS = """5<text class="displacy-token" fill="currentColor" text-anchor="start" y="{y}">6    <tspan class="displacy-word" fill="currentColor" x="{x}">{text}</tspan>7    <tspan class="displacy-tag" dy="2em" fill="currentColor" x="{x}">{tag}</tspan>8</text>9"""10 11TPL_DEP_SVG = """12<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" xml:lang="{lang}" id="{id}" class="displacy" width="{width}" height="{height}" direction="{dir}" style="max-width: none; height: {height}px; color: {color}; background: {bg}; font-family: {font}; direction: {dir}">{content}</svg>13"""14 15TPL_DEP_ARCS = """16<g class="displacy-arrow">17    <path class="displacy-arc" id="arrow-{id}-{i}" stroke-width="{stroke}px" d="{arc}" fill="none" stroke="red"/>18    <text dy="1.25em" style="font-size: 0.8em; letter-spacing: 1px">19        <textPath xlink:href="#arrow-{id}-{i}" class="displacy-label" startOffset="50%" side="{label_side}" fill="red" text-anchor="middle">{label}</textPath>20    </text>21    <path class="displacy-arrowhead" d="{head}" fill="red"/>22</g>23"""24 25 26def get_pil_text_size(text, font_size, font_name):27    font = ImageFont.truetype(font_name, font_size)28    size = font.getsize(text)29    return size30 31 32def render_arrow(33        label: str, start: int, end: int, direction: str, i: int34) -> str:35    """Render individual arrow.36 37    label (str): Dependency label.38    start (int): Index of start word.39    end (int): Index of end word.40    direction (str): Arrow direction, 'left' or 'right'.41    i (int): Unique ID, typically arrow index.42    RETURNS (str): Rendered SVG markup.43    """44 45    arc = get_arc(start + 10, 50, 5, end + 10)46    arrowhead = get_arrowhead(direction, start + 10, 50, end + 10)47    label_side = "right" if direction == "rtl" else "left"48    return TPL_DEP_ARCS.format(49        id=0,50        i=0,51        stroke=2,52        head=arrowhead,53        label=label,54        label_side=label_side,55        arc=arc,56    )57 58 59def get_arc(x_start: int, y: int, y_curve: int, x_end: int) -> str:60    """Render individual arc.61 62    x_start (int): X-coordinate of arrow start point.63    y (int): Y-coordinate of arrow start and end point.64    y_curve (int): Y-corrdinate of Cubic Bézier y_curve point.65    x_end (int): X-coordinate of arrow end point.66    RETURNS (str): Definition of the arc path ('d' attribute).67    """68    template = "M{x},{y} C{x},{c} {e},{c} {e},{y}"69    return template.format(x=x_start, y=y, c=y_curve, e=x_end)70 71 72def get_arrowhead(direction: str, x: int, y: int, end: int) -> str:73    """Render individual arrow head.74 75    direction (str): Arrow direction, 'left' or 'right'.76    x (int): X-coordinate of arrow start point.77    y (int): Y-coordinate of arrow start and end point.78    end (int): X-coordinate of arrow end point.79    RETURNS (str): Definition of the arrow head path ('d' attribute).80    """81    arrow_width = 682    if direction == "left":83        p1, p2, p3 = (x, x - arrow_width + 2, x + arrow_width - 2)84    else:85        p1, p2, p3 = (end, end + arrow_width - 2, end - arrow_width + 2)86    return f"M{p1},{y + 2} L{p2},{y - arrow_width} {p3},{y - arrow_width}"87 88 89def render_sentence_custom(unmatched_list: Dict, nlp):90    arcs_svg = []91    doc = nlp(unmatched_list["sentence"])92 93    x_value_counter = 1094    index_counter = 095    svg_words = []96    words_under_arc = []97    direction_current = "rtl"98 99    if unmatched_list["cur_word_index"] < unmatched_list["target_word_index"]:100        min_index = unmatched_list["cur_word_index"]101        max_index = unmatched_list["target_word_index"]102        direction_current = "left"103    else:104        max_index = unmatched_list["cur_word_index"]105        min_index = unmatched_list["target_word_index"]106    for i, token in enumerate(doc):107        word = str(token)108        word = word + " "109        pixel_x_length = get_pil_text_size(word, 16, 'arial.ttf')[0]110        svg_words.append(TPL_DEP_WORDS.format(text=word, tag="", x=x_value_counter, y=70))111        if min_index <= index_counter <= max_index:112            words_under_arc.append(x_value_counter)113            if index_counter < max_index - 1:114                x_value_counter += 50115        index_counter += 1116        x_value_counter += pixel_x_length + 4117 118    arcs_svg.append(render_arrow(unmatched_list['dep'], words_under_arc[0], words_under_arc[-1], direction_current, i))119 120    content = "".join(svg_words) + "".join(arcs_svg)121 122    full_svg = TPL_DEP_SVG.format(123        id=0,124        width=1200,  # 600125        height=75,  # 125126        color="#00000",127        bg="#ffffff",128        font="Arial",129        content=content,130        dir="ltr",131        lang="en",132    )133    return full_svg134