sembr 0.0.1__py3-none-any.whl

sembr/cli.py

@@ -0,0 +1,159 @@
+import os
+import sys
+import traceback
+
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+
+
+def parse_args():
+    import argparse
+    parser = argparse.ArgumentParser()
+    model_name = 'admko/sembr2023-bert-small'
+    parser.add_argument('-m', '--model-name', type=str, default=model_name)
+    parser.add_argument('-i', '--input-file', type=str, default=None)
+    parser.add_argument('-o', '--output-file', type=str, default=None)
+    parser.add_argument('-w', '--words-per-line', type=int, default=10)
+    parser.add_argument('-b', '--batch-size', type=int, default=8)
+    parser.add_argument('-d', '--overlap-divisor', type=int, default=8)
+    parser.add_argument(
+        '-f', '--predict-func', type=str,
+        choices=['argmax', 'breaks_first', 'logit_adjustment'],
+        default='argmax')
+    parser.add_argument('-t', '--tokens-per-line', type=int, default=10)
+    parser.add_argument('-s', '--server', type=str, default='127.0.0.1')
+    parser.add_argument('-l', '--listen', action='store_true')
+    parser.add_argument('-p', '--port', type=int, default=8384)
+    return parser.parse_args()
+
+
+def init(model_name):
+    import torch
+    from transformers import (AutoTokenizer, AutoModelForTokenClassification)
+    from .process import SemBrProcessor
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForTokenClassification.from_pretrained(model_name)
+    model.eval()
+    if torch.cuda.is_available():
+        model = model.to('cuda')
+    elif torch.backends.mps.is_available():
+        model = model.to('mps')
+    processor = SemBrProcessor()
+    return tokenizer, model, processor
+
+
+
+def start_server(port, tokenizer, model, processor):
+    from flask import Flask, request
+    app = Flask(__name__)
+    base_rv = {
+        'model': model.__class__.__name__,
+        'tokenizer': tokenizer.__class__.__name__,
+        'processor': processor.__class__.__name__,
+    }
+
+    @app.route('/check')
+    def check():
+        return {**base_rv, 'status': 'success'}
+
+    @app.route('/rewrap', methods=['POST'])
+    def rewrap():
+        from .inference import sembr
+        form = request.form
+        text = form['text']
+        kwargs = {
+            'batch_size': int(form.get('batch_size', 8)),
+            'predict_func': form.get('predict_func', 'argmax'),
+            'tokens_per_line': int(form.get('tokens_per_line', 10)),
+            'overlap_divisor': int(form.get('overlap_divisor', 8)),
+        }
+        try:
+            results = sembr(text, tokenizer, model, processor, **kwargs)
+            return {
+                **base_rv,
+                'status': 'success',
+                'text': results
+            }
+        except Exception as e:
+            return {
+                **base_rv,
+                'status': 'error',
+                'error': str(e),
+                'traceback': traceback.format_exc(),
+            }
+
+    app.run(port=port)
+
+
+def check_server(server, port):
+    import requests
+    from requests.exceptions import ConnectionError, ReadTimeout
+    if not server:
+        return False
+    try:
+        response = requests.get(f'http://{server}:{port}/check', timeout=0.3)
+    except (ConnectionError, ReadTimeout) as e:
+        return False
+    if response.status_code != 200:
+        return False
+    if response.json()['status'] != 'success':
+        return False
+    return True
+
+
+def rewrap_on_server(text, server, port, kwargs):
+    import requests
+    data = {
+        'text': text,
+        **kwargs,
+    }
+    try:
+        results = requests.post(
+            f'http://{server}:{port}/rewrap', data=data)
+    except Exception as e:
+        raise ValueError(f'Connection Error: {e}')
+    if results.status_code != 200:
+        raise ValueError(
+            f'Connection Error: {results.status_code}: {results.text}')
+    data = results.json()
+    if data['status'] != 'success':
+        raise ValueError(
+            f'Status: {data["status"]}\n'
+            f'Exception: {data["error"]}\n'
+            f'{data["traceback"]}')
+    return data['text']
+
+
+def main(args=None):
+    if args is None:
+        args = parse_args()
+    if args.listen:
+        tokenizer, model, processor = init(args.model_name)
+        return start_server(args.port, tokenizer, model, processor)
+    if args.input_file is not None:
+        with open(args.input_file, 'r', encoding='utf-8') as f:
+            text = f.read()
+    else:
+        text = sys.stdin.read()
+    kwargs = {
+        'batch_size': args.batch_size,
+        'predict_func': args.predict_func,
+        'tokens_per_line': args.tokens_per_line,
+        'overlap_divisor': args.overlap_divisor,
+    }
+    if check_server(args.server, args.port):
+        result = rewrap_on_server(text, args.server, args.port, kwargs)
+    else:
+        from .inference import sembr
+        tokenizer, model, processor = init(args.model_name)
+        result = sembr(text, tokenizer, model, processor, **kwargs)
+    if args.output_file is None:
+        print(result)
+        return
+    with open(args.output_file, 'w', encoding='utf-8') as f:
+        f.write(result)
+
+
+if __name__ == '__main__':
+    main(parse_args())

sembr/dataset.py

@@ -0,0 +1,93 @@
+import os
+import argparse
+import functools
+
+import datasets
+
+
+def _process_examples(
+    examples, processor, tokenizer, mode_names, max_indent, label2id
+):
+    examples['modes'] = [
+        [mode_names[i] for i in bm] for bm in examples['modes']]
+    transposed = [dict(zip(examples, c)) for c in zip(*examples.values())]
+    results = processor.tokenize_with_modes(tokenizer, transposed)
+    for r in results:
+        r['labels'] = labels = []
+        indents = []
+        for m, i in zip(r.pop('modes'), r.pop('indents')):
+            i = min(i, max_indent)
+            indents.append(i)
+            if m == 'off':
+                label = 'off'
+            else:
+                label = f'{m}-{i}'
+            labels.append(label2id[label])
+        r.pop('base_indent')
+    keys = ['input_ids', 'labels']
+    return {k: [d[k] for d in results] for k in keys}
+
+
+def chunk_examples(examples):
+    id_chunks, label_chunks = [], []
+    max_length = 512
+    overlap = int(max_length / 10)
+    for ids, labels in zip(examples.pop('input_ids'), examples.pop('labels')):
+        id_chunks += [
+            ids[i:i + max_length] for i in range(0, len(ids), overlap)]
+        label_chunks += [
+            labels[i:i + max_length] for i in range(0, len(labels), overlap)]
+    return {'input_ids': id_chunks, 'labels': label_chunks}
+
+
+def process_dataset(dataset, processor, tokenizer, max_indent, label2id):
+    mode_names = dataset.features['modes'].feature.names
+    removed_columns = [
+        'flat_lines', 'modes', 'mode_offsets', 'indents', 'base_indent']
+    process_examples = functools.partial(
+        _process_examples,
+        processor=processor, tokenizer=tokenizer, mode_names=mode_names,
+        max_indent=max_indent, label2id=label2id)
+    dataset = dataset.map(
+        process_examples, batched=True, remove_columns=removed_columns)
+    dataset = dataset.map(chunk_examples, batched=True)
+    return dataset
+
+
+def init_dataset():
+    dataset_file = os.path.join(os.path.dirname(__file__), 'sembr2023.py')
+    return datasets.load_dataset(dataset_file)
+
+def push_to_hub(dataset, hub_user, dataset_name, private):
+    hub_path = f'{hub_user}/{dataset_name}'
+    print(f'Pushing dataset to {hub_path}...')
+    dataset.push_to_hub(hub_path, private=private)
+
+
+def save_to_disk(dataset, save_dir, dataset_name):
+    save_path = os.path.join(save_dir, dataset_name)
+    print(f'Saving dataset to {save_path}...')
+    dataset.save_to_disk(save_path)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-n', '--dataset-name', type=str, default='sembr2023')
+    parser.add_argument('-s', '--save-dir', type=str, default='data')
+    parser.add_argument('-u', '--hub-user', type=str, default=None)
+    parser.add_argument('-p', '--private', action='store_true')
+    return parser.parse_args()
+
+
+def main(args=None):
+    if args is None:
+        args = parse_args()
+    dataset = init_dataset()
+    if args.save_dir is not None:
+        save_to_disk(dataset, args.save_dir, args.dataset_name)
+    if args.hub_user is not None:
+        push_to_hub(dataset, args.hub_user, args.dataset_name, args.private)
+
+
+if __name__ == '__main__':
+    main()

sembr/eval.py

@@ -0,0 +1,36 @@
+import glob
+
+import datasets
+import evaluate
+from transformers import AutoModelForTokenClassification, AutoTokenizer
+
+from .inference import inference
+from .process import SemBrProcessor
+
+
+def checkpoints():
+    return glob.glob("checkpoints/*/checkpoint-*")
+
+
+def eval_model(dataset, processor, checkpoint, metric):
+    model = AutoModelForTokenClassification.from_pretrained(checkpoint)
+    model = model.to('cuda')
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+    paras = dataset['test']['flat_lines']
+    results = inference(
+        paras, tokenizer, model, processor, batch_size=8, overlap_divisor=8)
+    generated = processor.generate(results, join=False)
+    return metric.compute(predictions=generated, references=paras)
+
+
+def main():
+    dataset = datasets.load_from_disk('./data/sembr2023')
+    processor = SemBrProcessor()
+    wer = evaluate.load('wer')
+    for checkpoint in checkpoints():
+        metrics = eval_model(dataset, processor, checkpoint, wer)
+        print(f'{checkpoint=}, {metrics=}')
+
+
+if __name__ == '__main__':
+    main()

sembr/inference.py

@@ -0,0 +1,119 @@
+import torch
+import numpy as np
+from tqdm import trange
+
+from transformers import DataCollatorForTokenClassification
+
+
+def _tiled_inference(model, collator, results, batch_size, overlap_divisor):
+    device = model.device
+    max_length = model.config.max_position_embeddings
+    overlap_length = int(max_length / overlap_divisor)
+    input_ids = [{'input_ids': r['input_ids']} for r in results]
+    num_paras = len(input_ids)
+    lengths = [len(i['input_ids']) for i in input_ids]
+    sorted_indices = sorted(
+        range(num_paras), key=lambda i: lengths[i], reverse=True)
+    logits = torch.zeros(
+        (num_paras, max(lengths), model.config.num_labels), device=device)
+    counts = torch.zeros(
+        (num_paras, max(lengths)), dtype=torch.long, device=device)
+    for b in trange(0, num_paras, batch_size):
+        bslice = slice(b, min(num_paras, b + batch_size))
+        bindices = sorted_indices[bslice]
+        binids = [input_ids[i] for i in bindices]
+        data = collator(binids, return_tensors='pt').to(device)
+        num_tokens = data['input_ids'].shape[1]
+        for i in range(0, num_tokens, max_length - overlap_length):
+            islice = slice(i, min(num_tokens, i + max_length))
+            inids = data['input_ids'][:, islice]
+            attns = data['attention_mask'][:, islice]
+            with torch.no_grad():
+                outputs = model(
+                    input_ids=inids, attention_mask=attns, return_dict=True)
+            logits[bindices, islice] += outputs.logits
+            counts[bindices, islice] += attns
+    attns = counts > 0
+    logits /= counts.unsqueeze(-1)
+    logits[~attns] = 0
+    return logits, attns
+
+
+def _format_labels(id2label, preds, attns, results):
+    modes, indents = [], []
+    for i, (p, a) in enumerate(zip(preds, attns)):
+        para_modes, para_indents = [], []
+        for name in [id2label[int(t)] for t in p[a]]:
+            if name == 'off':
+                mode, indent = 'off', 0
+            else:
+                mode, indent = name.split('-')
+            para_modes.append(mode)
+            para_indents.append(int(indent))
+        modes.append(para_modes)
+        indents.append(para_indents)
+    for r, m, i in zip(results, modes, indents):
+        r['modes'] = m
+        r['indents'] = i
+    return results
+
+
+def predict_argmax(logits, counts, tokens_per_line):
+    return logits.argmax(dim=2)
+
+
+def zero_runs(a, dim):
+    # Create an array that is 1 where a is 0, and pad each end with an extra 0.
+    pad = np.zeros_like(a.take([0], dim))
+    iszero = np.concatenate([pad, (a == 0).view(np.int8), pad], dim)
+    absdiff = np.abs(np.diff(iszero))
+    # Runs start and end where absdiff is 1.
+    return np.where(absdiff == 1)[0].reshape(-1, 2)
+
+
+def predict_logit_adjustment(logits, counts, tokens_per_line):
+    delta = 1.0
+    logits[:, :, 0] -= delta
+    logits[:, :, 1:] += delta / logits.shape[2]
+    return logits.argmax(dim=2)
+
+
+def predict_breaks_first(logits, counts, tokens_per_line):
+    delta = 0.0
+    off_logits = logits[:, :, 0] - delta
+    break_logits = logits[:, :, 1:] + delta / logits.shape[2]
+    breaks = (break_logits > off_logits.unsqueeze(-1)).any(2)
+    break_preds = 1 + break_logits.argmax(dim=2)
+    return torch.where(breaks, break_preds, torch.zeros_like(break_preds))
+
+
+PREDICT_FUNC_MAP = {
+    'argmax': predict_argmax,
+    'breaks_first': predict_breaks_first,
+    'logit_adjustment': predict_logit_adjustment,
+}
+
+
+def inference(
+    text, tokenizer, model, processor,
+    predict_func='argmax', tokens_per_line=10,
+    batch_size=8, overlap_divisor=8,
+):
+    collator = DataCollatorForTokenClassification(tokenizer, padding='longest')
+    results = processor(text, split=isinstance(text, str))
+    results = processor.tokenize_with_modes(tokenizer, results)
+    logits, counts = _tiled_inference(
+        model, collator, results, batch_size, overlap_divisor)
+    preds = PREDICT_FUNC_MAP[predict_func](logits, counts, tokens_per_line)
+    return _format_labels(model.config.id2label, preds, counts, results)
+
+
+def sembr(
+    text, tokenizer, model, processor,
+    predict_func='argmax', tokens_per_line=10,
+    batch_size=8, overlap_divisor=8,
+):
+    results = inference(
+        text, tokenizer, model, processor, predict_func, tokens_per_line,
+        batch_size, overlap_divisor)
+    return processor.generate(results, join=True)

sembr/process.py

@@ -0,0 +1,302 @@
+import re
+
+from transformers import AutoTokenizer
+
+
+class SemBrProcessor(object):
+    def __init__(self, spaces=4):
+        super().__init__()
+        self.spaces = spaces
+        self.replace_tokens = {
+            # r'\n(?:\s*\n)+': '[par]',
+            # '\t': '[indent]',
+            # ' ' * self.spaces: '[indent]',
+            '\\%': '[percent]',
+            '\n': '[newline]',
+        }
+        self.reverse_replace_tokens = {
+            v: k for k, v in self.replace_tokens.items()}
+
+    def prepare_tokenizer(self, tokenizer):
+        tokenizer.add_tokens(list(self.replace_tokens.values()))
+
+    def _process_specials(self, lines):
+        for k, v in self.replace_tokens.items():
+            # lines = [re.sub(k, v, l) for l in lines]
+            lines = [l.replace(k, v) for l in lines]
+        return lines
+
+    def _process_indents(self, lines):
+        nlines = []
+        indents = []
+        # get indent levels
+        for line in lines:
+            indent = 0
+            for c in line:
+                if c == ' ':
+                    indent += 1
+                elif c == '\t':
+                    indent += self.spaces
+                else:
+                    break
+            indent_level = int(indent / self.spaces)
+            nlines.append(line[indent_level * self.spaces:].rstrip())
+            indents.append(indent_level)
+        return nlines, indents
+
+    def _process_comments(self, lines, indents):
+        # normalize comments, ['xxx % comment'] -> ['xxx', '% comment']
+        nclines = []
+        ncindents = []
+        for line, indent in zip(lines, indents):
+            if '%' in line:
+                normal, *comment = line.split('%')
+                comment = '%'.join(comment).strip()
+                if normal.strip():
+                    if comment:
+                        nclines += [normal, f'%{comment}']
+                        ncindents += [indent, indent]
+                        continue
+                    line = f'{normal}%'
+            nclines.append(line)
+            ncindents.append(indent)
+        return nclines, ncindents
+
+    def _process_modes(self, lines):
+        new_lines = []
+        modes = []
+        prev_status = 'start'
+        for line in lines:
+            if line.startswith('%'):
+                status = 'comment'
+            elif line.endswith('%'):
+                status = 'percent'
+                line = line.rstrip('%')
+            else:
+                status = 'normal'
+            match (prev_status, status):
+                case ('start', _):
+                    pass
+                case ('normal', _):
+                    modes.append('space')
+                case ('percent', _):
+                    modes.append('nospace')
+                case ('comment', 'normal'):
+                    modes.append('break')
+                case ('comment', 'percent'):
+                    modes.append('break')
+                case ('comment', 'comment'):
+                    modes.append('comment')
+                case (_, 'comment'):
+                    modes.append('comment')
+                case _:
+                    raise ValueError(
+                        'Unknown status transition: '
+                        f'{prev_status} -> {status}.')
+            new_lines.append(line)
+            prev_status = status
+        # last transition always force a break
+        modes.append('break')
+        return new_lines, modes
+
+    def _flatten_with_modes(self, lines, modes):
+        in_comment = 0
+        flat_lines, flat_modes, offsets = [], [], []
+        prev_len = flat_len = 0
+        for line, mode in zip(lines, modes):
+            if in_comment >= 1:
+                line = re.sub(r'^\s*%', '', line)
+            if mode == 'break':
+                in_comment = 0
+                line = f'{line}[newline]'
+                mode = 'off'
+            elif mode == 'comment':
+                in_comment += 1
+                mode = 'space'
+            elif mode == 'space':
+                line = f'{line} '
+            elif mode == 'nospace':
+                pass
+            else:
+                raise ValueError(f'Unknown mode: {mode}.')
+            flat_lines.append(line)
+            flat_modes.append(mode)
+            prev_len = flat_len
+            flat_len += len(line)
+            offsets.append((prev_len, flat_len))
+        return ''.join(flat_lines), flat_modes, offsets
+
+    def _process_paragraph(self, text):
+        lines = text.split('\n')
+        lines = self._process_specials(lines)
+        lines, indents = self._process_indents(lines)
+        base_indent = min(indents)
+        indents = [i - base_indent for i in indents]
+        lines, indents = self._process_comments(lines, indents)
+        lines, modes = self._process_modes(lines)
+        flat_lines, modes, mode_offsets = self._flatten_with_modes(lines, modes)
+        return {
+            'flat_lines': flat_lines,
+            'modes': modes,
+            'mode_offsets': mode_offsets,
+            'indents': indents,
+            'base_indent': base_indent,
+        }
+
+    def _tokenize_with_modes(
+        self, tokenizer, text, line_modes, line_mode_offsets, line_indents
+    ):
+        enc = tokenizer(text, return_offsets_mapping=True)
+        words, modes, indents = [], [], []
+        pos = mode_idx = 0
+        # fill empty words in offset mapping
+        offset_mapping = []
+        for start, end in enc.offset_mapping:
+            offset_mapping.append((min(start, pos), end))
+            pos = end
+        pos = 0
+        input_ids = []
+        for tid, (start, end) in zip(enc.input_ids, offset_mapping):
+            if pos >= len(text):
+                break
+            mode_offset = line_mode_offsets[mode_idx][1]
+            word = text[pos:end]
+            input_ids.append(tid)
+            words.append(word)
+            indents.append(line_indents[mode_idx])
+            pos = max(pos, end)
+            if mode_offset >= end:
+                modes.append('off')
+                continue
+            mode = line_modes[mode_idx]
+            modes.append(mode)
+            mode_idx += 1
+            # current word is on a new line
+            indents[-1] = line_indents[mode_idx]
+        return input_ids, words, modes, indents
+
+    def tokenize_with_modes(self, tokenizer, results):
+        self.prepare_tokenizer(tokenizer)
+        new_results = []
+        for r in results:
+            flat_lines = r['flat_lines']
+            modes = r['modes']
+            mode_offsets = r['mode_offsets']
+            indents = r['indents']
+            base_indent = r['base_indent']
+            input_ids, words, modes, indents = self._tokenize_with_modes(
+                tokenizer, flat_lines, modes, mode_offsets, indents)
+            tokenized = {
+                'input_ids': input_ids,
+                'words': words,
+                'modes': modes,
+                'indents': indents,
+                'base_indent': base_indent,
+            }
+            keys = ['input_ids', 'words', 'modes', 'indents']
+            if len(set(len(tokenized[k]) for k in keys)) != 1:
+                len_dict = {k: len(tokenized[k]) for k in keys}
+                raise ValueError(
+                    f'Lengths do not match. Found: {len_dict}.')
+            new_results.append(tokenized)
+        return new_results
+
+    def __call__(self, text, split=True):
+        if split:
+            text = re.split(r'\n(?:\s*\n)+', text)
+        elif isinstance(text, str):
+            raise ValueError(
+                'Text must be a list of strings if split=True.')
+        paras = []
+        for p in text:
+            if not p.strip():
+                continue
+            paras.append(self._process_paragraph(p))
+        return paras
+
+    def _replace_newlines(self, words, modes, indents):
+        new_words, new_modes, new_indents = [], [], []
+        next_mode = None
+        for word, mode, indent in zip(words, modes, indents):
+            if word == '[newline]':
+                next_mode = 'break'
+                continue
+            if next_mode:
+                # if mode != 'off':
+                #     raise ValueError(
+                #         f'Cannot set mode {next_mode} '
+                #         f'when mode is {mode}.')
+                mode = next_mode
+                next_mode = None
+            new_words.append(word)
+            new_modes.append(mode)
+            new_indents.append(indent)
+        return new_words, new_modes, new_indents
+
+    def _generate_lines(self, words, modes, indents):
+        lbs = [
+            (o, m) for o, m in enumerate(modes)
+            if m in ('space', 'nospace', 'break')]
+        if not lbs or lbs[-1][0] < len(words):
+            lbs.append((len(words), 'space'))
+        lines, line_indents  = [], []
+        pos = in_comment = 0
+        for o, m in lbs:
+            line = ''.join(words[pos:o]).strip()
+            if line.startswith('%'):
+                in_comment = 1
+            if m == 'nospace':
+                line = f'{line}%'
+            if m in ('space', 'break'):
+                if in_comment > 1:
+                    line = f'% {line}'
+                if in_comment:
+                    in_comment += 1
+                if m == 'break':
+                    in_comment = 0
+            lines.append(line)
+            line_indents.append(indents[pos:o])
+            pos = o
+        # line_indents = [Counter(l).most_common(1)[0][0] for l in line_indents]
+        line_indents = [l[0] for l in line_indents]
+        return lines, line_indents
+
+    def _indent_lines(self, lines, indents, base_indent):
+        spaces = ' ' * self.spaces
+        return [
+            f'{spaces * (i + base_indent)}{l}'
+            for i, l in zip(indents, lines)]
+
+    def _generate_paragraph(self, processed):
+        words = processed['words']
+        modes = processed['modes']
+        indents = processed['indents']
+        base_indent = processed['base_indent']
+        words, modes, indents = self._replace_newlines(words, modes, indents)
+        lines, indents = self._generate_lines(words, modes, indents)
+        lines = self._indent_lines(lines, indents, base_indent)
+        text = '\n'.join(lines)
+        for k, v in self.reverse_replace_tokens.items():
+            text = text.replace(k, v)
+        return text
+
+    def generate(self, paragraphs, join=True):
+        paragraphs = [self._generate_paragraph(p) for p in paragraphs]
+        if join:
+            return '\n\n'.join(paragraphs)
+        return paragraphs
+
+
+if __name__ == '__main__':
+    # test = open('./data/test/mair.tex', 'r').read()
+    test = open('./data/example.tex', 'r').read()
+    processor = SemBrProcessor()
+    tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
+    results = processor(test)
+    results = processor.tokenize_with_modes(tokenizer, results)
+    print('--- Processed ---')
+    print(processor.generate(results))
+    for r in results:
+        r['modes'] = ['off' if m != 'break' else m for m in r['modes']]
+    print('--- Flattened ---')
+    print(processor.generate(results))

sembr/sembr2023.py

@@ -0,0 +1,60 @@
+import os
+import glob
+
+import datasets
+
+from .process import SemBrProcessor
+
+
+logger = datasets.logging.get_logger(__name__)
+
+MAX_INDENT = 10
+
+
+class SemBr2023(datasets.GeneratorBasedBuilder):
+    BUILDER_CONFIGS = [
+        datasets.BuilderConfig(
+            name='sembr2023',
+            version=datasets.Version('1.0.0'),
+            description='SemBr2023 dataset'),
+    ]
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.processor = SemBrProcessor()
+
+    def _info(self):
+        modes = ['off', 'space', 'nospace']
+        indents = [str(i) for i in range(MAX_INDENT + 1)]
+        return datasets.DatasetInfo(
+            features=datasets.Features({
+                'flat_lines': datasets.Value('string'),
+                'modes': datasets.Sequence(
+                    datasets.features.ClassLabel(names=modes)),
+                'mode_offsets': datasets.Sequence(
+                    datasets.Sequence(datasets.Value('int32'))),
+                'indents': datasets.Sequence(
+                    datasets.features.ClassLabel(names=indents)),
+                'base_indent': datasets.Value('int32'),
+            })
+        )
+
+    def _split_generators(self, dl_manager):
+        return [
+            datasets.SplitGenerator(
+                name=datasets.Split.TRAIN,
+                gen_kwargs={'root': './data/raw/train/'}),
+            datasets.SplitGenerator(
+                name=datasets.Split.TEST,
+                gen_kwargs={'root': './data/raw/test/'}),
+        ]
+
+    def _generate_examples(self, root):
+        eid = 0
+        for path in glob.glob(os.path.join(root, '*')):
+            logger.info(f'Generating examples from {path!r}...')
+            with open(path, 'r', encoding='utf-8') as f:
+                text = f.read()
+            for p in self.processor(text):
+                yield eid, p
+                eid += 1

sembr/train.py

@@ -0,0 +1,128 @@
+import os
+import sys
+import argparse
+
+import datasets
+from transformers import (
+    AutoTokenizer, AutoModelForTokenClassification,
+    TrainingArguments, Trainer, DataCollatorForTokenClassification)
+
+from .process import SemBrProcessor
+from .dataset import process_dataset
+from .utils import compute_metrics
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('model', type=str)
+    parser.add_argument('-dn', '--dataset-name', type=str, default='admko/sembr2023')
+    parser.add_argument('-lr', '--learning-rate', type=float, default=1e-5)
+    parser.add_argument('-tb', '--train-batch-size', type=int, default=64)
+    parser.add_argument('-eb', '--eval-batch-size', type=int, default=128)
+    parser.add_argument('-mi', '--max-indent', type=int, default=3)
+    parser.add_argument('-hu', '--hub-user', type=str, default=None)
+    parser.add_argument('-ms', '--max-steps', type=int, default=5000)
+    parser.add_argument('-es', '--eval-steps', type=int, default=10)
+    parser.add_argument('-ss', '--save-steps', type=int, default=100)
+    parser.add_argument('-rt', '--report-to', type=str, default='all')
+    parser.add_argument('-d', '--debug', action='store_true')
+    args = parser.parse_args()
+    if args.debug:
+        import debugpy
+        debugpy.listen(5678)
+        print('Waiting for debugger...')
+        debugpy.wait_for_client()
+        args.report_to = 'none'
+        args.hub_user = None
+    return args
+
+
+class DataCollatorForTokenClassificationWithTruncation(
+    DataCollatorForTokenClassification
+):
+    def __init__(self, tokenizer, max_length=512, **kwargs):
+        super().__init__(tokenizer, **kwargs)
+        self.max_length = max_length
+
+    def __call__(self, features, return_tensors=None):
+        truncated_features = []
+        for f in features:
+            truncated_features.append(
+                {k: v[:self.max_length] for k, v in f.items()})
+        return super().__call__(truncated_features, return_tensors)
+
+
+def init_dataset(args, label2id, max_length):
+    dataset = datasets.load_dataset(args.dataset_name)
+    processor = SemBrProcessor()
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    processor.prepare_tokenizer(tokenizer)
+    train_dataset = process_dataset(
+        dataset['train'], processor, tokenizer, args.max_indent, label2id)
+    test_dataset = process_dataset(
+        dataset['test'], processor, tokenizer, args.max_indent, label2id)
+    print(f'{len(train_dataset)=}, {len(test_dataset)=}')
+    collator = DataCollatorForTokenClassificationWithTruncation(
+        tokenizer, padding='max_length', max_length=max_length)
+    return train_dataset, test_dataset, tokenizer, collator
+
+
+def init_model(model_name, max_indent):
+    label_names = ['off'] + [
+        f'{m}-{i}' for m in ['space', 'nospace']
+        for i in range(max_indent + 1)]
+    id2label = {i: l for i, l in enumerate(label_names)}
+    label2id = {l: i for i, l in enumerate(label_names)}
+    model = AutoModelForTokenClassification.from_pretrained(
+        model_name, ignore_mismatched_sizes=True,
+        num_labels=len(id2label), id2label=id2label, label2id=label2id)
+    return model
+
+
+def main(args):
+    model = init_model(args.model, args.max_indent)
+    max_length = model.config.max_position_embeddings
+    train_dataset, test_dataset, tokenizer, collator = \
+        init_dataset(args, model.config.label2id, max_length)
+    model.config.__dict__['max_indent'] = args.max_indent
+    model.resize_token_embeddings(len(tokenizer))
+    model_name = args.model.split('/')[-1]
+    run_name = f'sembr2023-{model_name}'
+    training_args = TrainingArguments(
+        output_dir=f'checkpoints/{run_name}',
+        run_name=run_name,
+        report_to=args.report_to,
+        learning_rate=args.learning_rate,
+        lr_scheduler_type='cosine',
+        per_device_train_batch_size=args.train_batch_size,
+        per_device_eval_batch_size=args.eval_batch_size,
+        weight_decay=1e-5,
+        evaluation_strategy='steps',
+        max_steps=args.max_steps,
+        eval_steps=args.eval_steps,
+        save_strategy='steps',
+        save_steps=args.save_steps,
+        save_total_limit=1,
+        metric_for_best_model='f1',
+        load_best_model_at_end=True,
+        logging_steps=1,
+        push_to_hub=args.hub_user is not None,
+        hub_strategy='end',
+        hub_model_id=f'{args.hub_user}/{run_name}',
+    )
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=test_dataset,
+        tokenizer=tokenizer,
+        data_collator=collator,
+        compute_metrics=compute_metrics,
+    )
+    trainer.train()
+    if args.hub_user is not None:
+        trainer.push_to_hub()
+
+
+if __name__ == '__main__':
+    main(parse_args())

sembr/utils.py

@@ -0,0 +1,44 @@
+import torch
+import numpy as np
+
+
+def binary_metrics(fp, tp, fn, tn=None, prefix=None):
+    def safe_div(a, b):
+        return a / b if b > 0 else 0
+    result = {
+        'precision': safe_div(tp, tp + fp),
+        'recall': safe_div(tp, tp + fn),
+        'f1': safe_div(2 * tp, 2 * tp + fp + fn),
+        'iou': safe_div(tp, tp + fp + fn),
+    }
+    if tn is not None:
+        accuracies = {
+            'accuracy': safe_div(tp + tn, tn + fn + tp + fp),
+            'balanced_accuracy': (
+                safe_div(tp, tp + fn) + safe_div(tn, tn + fp)) / 2,
+        }
+        result.update(accuracies)
+    if prefix is not None:
+        result = {f'{prefix}_{k}': v for k, v in result.items()}
+    return result
+
+
+def compute_metrics(result):
+    logits, labels = result
+    preds = np.argmax(logits, axis=2)
+    idx = labels != -100
+    logits, preds, labels = logits[idx], preds[idx], labels[idx]
+    loss = torch.nn.functional.cross_entropy(
+        torch.Tensor(logits), torch.LongTensor(labels), reduction='mean')
+    preds_set = set(preds.nonzero()[0])
+    labels_set = set(labels.nonzero()[0])
+    fp = len(preds_set - labels_set)
+    tp = len(labels_set & preds_set)
+    fn = len(labels_set - preds_set)
+    tn = len(set(range(len(preds))) - (labels_set | preds_set))
+    metrics = {
+        **binary_metrics(fp, tp, fn, tn),
+        'overall_accuracy': (preds == labels).mean(),
+        'loss': loss,
+    }
+    return metrics

sembr-0.0.1.dist-info/LICENSE.txt

@@ -0,0 +1,33 @@
+The MIT License (MIT)
+
+Copyright (c) 2023 Xitong Gao
+
+Permission is hereby granted,
+free of charge,
+to any person obtaining a copy of this software
+and associated documentation files (the "Software"),
+to deal in the Software without restriction,
+including without limitation
+the rights to use, copy, modify, merge,
+publish, distribute, sublicense,
+and/or sell copies of the Software,
+and to permit persons to whom the Software
+is furnished to do so,
+subject to the following conditions:
+
+The above copyright notice
+and this permission notice
+shall be included in all copies
+or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS",
+WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED,
+INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
+BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY,
+WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE
+OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

sembr-0.0.1.dist-info/METADATA

@@ -0,0 +1,256 @@
+Metadata-Version: 2.1
+Name: sembr
+Version: 0.0.1
+Summary: A semantic linebreaker powered by transformers
+Author: admk
+Project-URL: Homepage, https://github.com/admk/sembr
+Project-URL: Issues, https://github.com/admk/sembr/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+Requires-Dist: transformers
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: tqdm
+Requires-Dist: requests
+Requires-Dist: flask
+
+# Semantic Line Breaker (SemBr)
+
+```
+> When writing text
+> with a compatible markup language,
+> add a line break
+> after each substantial unit of thought.
+```
+
+## What is SemBr?
+
+SemBr is a tool
+that breaks lines in a text file
+at semantic boundaries.
+
+### Installation
+
+SemBr is available as a Python package
+on PyPI.
+To install it,
+simply run the following command
+in your terminal,
+assuming that you have Python 3.10 or later installed:
+```shell
+pip install sembr
+```
+
+### Supported Platforms
+
+SemBr is supported on Linux, Mac and Windows.
+On machines with CUDA devices,
+or on Apple Silicon Macs,
+SemBr will use the GPU / Apple Neural Engine
+to accelerate inference.
+
+### Usage
+
+To use SemBr,
+run the following command
+in your terminal:
+```shell
+sembr -i <input_file> -o <output_file>
+```
+where `<input_file>` is the path to the input file.
+
+Alternatively,
+you can pipe the input
+into `sembr`,
+and the output can also be printed
+to the terminal:
+```shell
+cat <input_file> | sembr
+```
+This is especially useful
+if you want to use SemBr
+with clipboard managers,
+for instance, on a Mac:
+```shell
+pbpaste | sembr | pbcopy
+```
+
+Additionally,
+you can specify the following options
+to customize the behavior of SemBr:
+* `-m <model_name>`: The name of the Hugging Face model to use.
+  The default is `admko/sembr2023-bert-small`.
+* `-l`: Serves the SemBr API on a local server.
+  Each instance of `sembr` run
+  will detect if the API is accessible,
+  and if not it will run the model on its own.
+  This option is useful
+  to avoid the time taken to initialize the model
+  by keeping it in memory in a separate process.
+* `-p <port>`: The port to serve the SemBr API on.
+  The default is `8384`.
+* `-s <ip>`: The IP address to serve the SemBr API on.
+  The default is `127.0.0.1`.
+
+## What are Semantic Line Breaks?
+
+[Semantic Line Breaks](https://sembr.org)
+or [
+    Semantic Linefeeds
+](https://rhodesmill.org/brandon/2012/one-sentence-per-line/)
+describe a set of conventions
+for using insensitive vertical whitespace
+to structure prose along semantic boundaries.
+
+## Why use Semantic Line Breaks?
+
+Semantic Line Breaks has the following advantages:
+
+* Breaking lines by splitting clauses
+  reflects the logical, grammatical and semantic structure
+  of the text.
+
+* It enhances the ease of editing and version control
+  for a text file.
+  Merge conflicts are less likely to occur
+  when small changes are made,
+  and the changes are easier to identify.
+
+* Documents written with semantic line breaks
+  are easier to navigate and edit
+  with Vim and other text editors
+  that use Vim keybindings.
+
+* Semantic line breaks
+  are invisible to readers.
+  The final rendered output
+  shows no changes to the source text.
+
+## Why SemBr?
+
+Converting existing text not written
+with semantic line breaks
+takes a long time to do it manually,
+and it is surprisingly difficult
+to do it automatically
+with rule-based methods.
+
+### Challenges of rule-based methods
+
+Rule-based heuristics do not work well
+with the actual semantic structure of the text,
+often leading to incorrect semantic boundaries.
+Moreover,
+semantic boundaries are hierarchical and nested,
+and a rule-based approach
+cannot capture this structure.
+A semantic line break
+may occur after a dependent clause,
+but not all clauses should be broken into lines.
+For examples:
+
+* A rule that breaks lines at punctuation marks
+  will not work well
+  with sentences that contain
+  periods in abbreviations or mathematical expressions.
+
+* For example,
+  "I like to eat apples and oranges
+  because they are healthy."
+  should be broken into lines as follows:
+  ```
+  I like to eat apples and oranges
+  because they are healthy.
+  ```
+  rather than:
+  ```
+  I like to eat apples
+  and oranges because they are healthy.
+  ```
+
+For this reason,
+I have created SemBr,
+which uses finetuned Transformer models
+to predict line breaks at semantic boundaries.
+
+
+## How does SemBr work?
+
+SemBr uses a Transformer model
+to predict line breaks
+at semantic boundaries.
+A small dataset of text
+with semantic line breaks was created
+from my existing LaTeX documents.
+The dataset was split into training
+(46,295 lines, 170,681 words and 1,492,952 characters)
+and test
+(2,187 lines, 7,564 words and 72,231 characters)
+datasets.
+
+The data was prepared
+by extracting line breaks and indent levels
+from the files,
+and then converting the files
+into strings of paragraphs
+with line breaks removed.
+The data can then be tokenized
+using the tokenizer
+and converted into a dataset
+with tokens,
+where each token has a label
+denoting:
+* no line break (label = 0), or
+* a line break
+  that adds a space in LaTeX documents
+  at the token with an indent level
+  (label in [0, 1, 2, ..., MAX_INDENT]), or
+* a line break that adds no space
+  (label in [MAX_INDENT + 1, MAX_INDENT + 2, ..., 2 * MAX_INDENT]).
+
+The pretrained masked language model
+is then finetuned as a token classifier
+on the training dataset
+to predict the labels
+of the tokens.
+We save the model
+with the best F1 score
+on correctly predicting line breaks of any kind
+on the test set.
+The finetuning logs
+for all models including the following
+can be found
+on this [WandB](https://wandb.ai/admko/sembr2023) report:
+* [`distilbert-base-uncased`](https://huggingface.co/distilbert-base-uncased).
+* [`distilbert-base-cased`](https://huggingface.co/distilbert-base-cased).
+* [`distilbert-base-uncased-finetuned-sst-2-english`](https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english).
+* [`prajjwal1/bert-tiny`](https://huggingface.co/prajjwal1/bert-tiny).
+* [`prajjwal1/bert-small`](https://huggingface.co/prajjwal1/bert-small).
+
+
+## Performance
+
+Current inference speed
+on an M2 Macbook Pro
+is about 1,500 characters per second
+on `bert-small`,
+the memory usage is about 1.70 GB.
+
+
+## Improvements and TODOs
+
+* [ ] Support natural languages other than English.
+* [ ] Support other markup languages
+      such as Markdown.
+* [ ] Some lines are too long
+      without a line break.
+      The inference algorithm
+      can be improved to penalize long lines.
+* [ ] Performance benchmarking.
+* [ ] Improve inference speed.
+* [ ] Reduce memory usage.
+* [ ] Improve indent level prediction.

sembr-0.0.1.dist-info/RECORD

@@ -0,0 +1,15 @@
+sembr/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sembr/cli.py,sha256=3uP01_v3-MRhwL3_46u1aG6p2BXq46DIE8SPA35Opu8,5179
+sembr/dataset.py,sha256=FZqXwfrXY8veN7YNI5hu0FM3fCNcxNBWzEsK8XySR9U,3177
+sembr/eval.py,sha256=fWIMlcRapnD3z0wj8ejvyaPBK4lx06n5bDke9oB4rHE,1066
+sembr/inference.py,sha256=80bWRBa1FSMfN30SN72ii9eboTuO9DHXyzGdBJjEhqQ,4361
+sembr/process.py,sha256=QgP2LbsgS2lboNvtKquiyWPhw3DSdPguB6WEg3-cRJ4,11017
+sembr/sembr2023.py,sha256=tEdZR6IThAML69gJM0ZtWi6mBIh6xTSGFnz8mE2xgGI,1878
+sembr/train.py,sha256=BpBybR8s0m8kLWEhOy5XOFNRffg6hqbddcMRYJzUnQw,4791
+sembr/utils.py,sha256=dp4xY_OME3P7aETYS16Xn55lEfqVg3K871Un7jllu_4,1436
+sembr-0.0.1.dist-info/LICENSE.txt,sha256=ox_0dcgQRQpWRLa3ZuFBzCoAbnCoKxFRMHnQPh7Mue0,1077
+sembr-0.0.1.dist-info/METADATA,sha256=cCgJMS2QtC56V0-iLcyZ50y-XFW8I4J2Ib6EukSzBPA,6775
+sembr-0.0.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+sembr-0.0.1.dist-info/entry_points.txt,sha256=TPj7T_mm1qglnW5yf8JWQjEY8BaAb3CWyH8TpFE1u8U,41
+sembr-0.0.1.dist-info/top_level.txt,sha256=q-9hKQBeQBW_3ERrGmDUlJjcaZla-VxX9PTK13v-u9E,6
+sembr-0.0.1.dist-info/RECORD,,

sembr-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.42.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

sembr-0.0.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+sembr = sembr.cli:main

sembr-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+sembr