newline-restoration-distilroberta

Newline Restoration, DistilRoBERTa

This model fixes newline placement in English text. Paste text whose line breaks were lost (a PDF copy, an OCR dump, a scraped page) and the model predicts where each line break belongs.

It labels every gap between two words as O (keep a space) or NEWLINE (start a new line).

Training

• Base model: distilroberta-base (6 layers, about 82M parameters).
• Data: 12k arXiv papers (neuralwork/arxiver, markdown cleaned to keep headings and bullets) mixed with 9k Wikipedia articles. Papers supply the headings, section numbers, and bullets the model infers on; Wikipedia keeps general prose in the mix.
• Objective: self-supervised. Well-formatted text already carries the line breaks, so training strips them and records each word's following whitespace as the label.
• Schedule: 4 epochs, learning rate 5e-5.

Threshold

Inference thresholds P(NEWLINE) at 0.25, stored in inference_config.json. We picked that value on validation with an F-beta score (beta 1.5) that favors recall, because a missed break is the visible failure.

Results (held-out test, 7.48M word boundaries)

Inference

The companion service rejoins words that a hard wrap split (que + ries becomes queries), slides a 200-word window and averages overlapping break probabilities, forces a break before bullets and section numbers, then breaks after a word once its probability clears the threshold.

from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch

repo = "iampoppyxx/newline-restoration-distilroberta"
tok = AutoTokenizer.from_pretrained(repo, add_prefix_space=True)
model = AutoModelForTokenClassification.from_pretrained(repo).eval()

@torch.inference_mode()
def restore(text, threshold=0.25):
    words = text.split()
    enc = tok(words, is_split_into_words=True, return_tensors="pt",
              truncation=True, max_length=256)
    probs = model(**enc).logits[0].softmax(-1)[:, 1]
    out, seen = [], set()
    for pos, wid in enumerate(enc.word_ids()):
        if wid is None or wid in seen:
            continue
        seen.add(wid)
        newline = wid < len(words) - 1 and probs[pos] >= threshold
        out.append(words[wid] + ("\n" if newline else " "))
    return "".join(out).strip()

print(restore("3.2.3 Applications of Attention in our Model The Transformer uses "
              "multi-head attention in three different ways."))