sobir-hf/tajik-text-segmentation

1import os2import re3 4 5def parse_annotated_text(text):6    # Regular expression pattern to parse YEDDA format7    yedda_pattern = re.compile(r'(\[\@(.*?)\#(\w+)\*\])', re.DOTALL)8    9    # This variable will keep track of the number of characters removed10    chars_removed = 011 12    # This will store the spans of the entities in the original text13    spans_in_original_text = []14 15    # Buffer to store content without annotations16    buffer = []17 18    # Index to track last end position19    last_end = 020 21    # Store labels22    labels = []23 24    # Loop through each match25    for match in yedda_pattern.finditer(text):26        # The entire match27        full_match = match.group(0)28        # Capture group 2 (entity)29        entity = match.group(2)30        # Capture group 2 (label)31        label = match.group(3)32        # Start position of the match in the modified string33        start = match.start()34        # End position of the match in the modified string35        end = match.end()36 37        labels.append(label)38 39        # Append the text before the annotation to the buffer40        buffer.append(text[last_end:start])41 42        # Remove trailing spaces after entity43        entity = entity.rstrip()44        buffer.append(entity)45 46        # Calculate the start and end spans in the original text47        original_start = start - chars_removed48        original_end = original_start + len(entity)49 50        assert original_end > original_start, text51        52        # Store the spans53        spans_in_original_text.append((original_start, original_end))54        55        # update the chars_removed counter56        chars_removed += len(full_match) - len(entity)57 58        # Update last_end59        last_end = end60 61    # Append remaining content after the last match62    buffer.append(text[last_end:])63 64    # Join buffer parts to get content without annotations65    content_without_annotations = "".join(buffer)66 67    return {68        'text': content_without_annotations,69        'spans': spans_in_original_text,70        'labels': labels71    }72 73 74def preprocess_text(text: str):75    # remove extra spaces76    text = text.strip()77    text = re.sub(r'\n+', '\n', text)78    text = re.sub(r' +', ' ', text)79    text = text.replace(' \n', '\n')80    text = text.replace('\n ', '\n')81    return text82 83 84def load_yedda_annotations(directory):85 86    # List to store all the annotations from all files87    all_annotations = []88    89    # Iterate through each file in the given directory90    for filename in os.listdir(directory):91        # Check if the file has the '.ann' extension92        if filename.endswith(".ann"):93            # Construct the full file path94            file_path = os.path.join(directory, filename)95            96            # Open and read the file97            with open(file_path, 'r', encoding='utf-8') as file:98                content = file.read()99            100            # Preprocess text101            content = preprocess_text(content)102            103            parsed = parse_annotated_text(content)104            file_annotations = {105                'file': filename, 106                'annotated_text': content, 107                'text': parsed['text'],108                'spans': parsed['spans'],109                'labels': parsed['labels'],110            }111            all_annotations.append(file_annotations)112    113    return all_annotations114 115 116def convert_to_ann(annotatations):117    text = annotatations['text']118    buffer = []119    i = 0120    for (j_start, j_end), label in zip(annotatations['spans'], annotatations['labels']):121 122        buffer += text[i:j_start]123        buffer += [f'[@{text[j_start:j_end]}#{label}*]']124        i = j_end125 126    buffer += [text[i:]]127 128    return ''.join(buffer)129 130 131if __name__ == '__main__':132    133    directory_path = 'annotations'  # The directory containing .ann files134    annotations = load_yedda_annotations(directory_path)135 136    counter = 0137    for file_annotation in annotations:138        counter += len(file_annotation['labels'])139        print('File:', file_annotation['file'])140        print('Text[:100]:', repr(file_annotation['text'][:100]))141        print('Number of labels:', len(file_annotation['labels']))142        assert len(file_annotation['labels']) == len(file_annotation['spans'])143        print('Average labeled sentence length:', sum(end-start for start,end in file_annotation['spans']) / len(file_annotation['spans']))144        print('--------------------------------')145    146    print('Total number of files:', len(annotations))147    print('Total label count:', counter)148 149