joelbarmettler/md-reheader
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: joelbarmettler/md-reheader
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-06-13 10:17:16 +08:00
commit 41ed79e963
8 changed files with 494 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
.gitattributes vendored Normal file
View File

@@ -0,0 +1,36 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.npy filter=lfs diff=lfs merge=lfs -text
*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.pickle filter=lfs diff=lfs merge=lfs -text
*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
tokenizer.json filter=lfs diff=lfs merge=lfs -text

269
README.md Normal file
View File

@@ -0,0 +1,269 @@
---
license: apache-2.0
language:
- en
tags:
- markdown
- document-structure
- heading-prediction
- pdf-to-markdown
- fine-tuned
base_model: Qwen/Qwen3-0.6B
datasets:
- joelbarmettler/md-reheader-dataset
metrics:
- accuracy
- exact_match
pipeline_tag: text-generation
library_name: transformers
---
<div align="center">
<h1 align="center" style="font-size: 32px">md-reheader</h1>
<p align="center"><strong>Restore heading hierarchy in markdown documents with a fine-tuned 0.6B LLM.</strong></p>
<p align="center">
<a href="https://pypi.org/project/md-reheader/"><img src="https://img.shields.io/pypi/v/md-reheader?color=blue&label=PyPI" alt="PyPI"></a>
<a href="https://www.python.org/"><img src="https://img.shields.io/badge/python-3.12%2B-blue" alt="Python 3.12+"></a>
<a href="https://www.apache.org/licenses/LICENSE-2.0"><img src="https://img.shields.io/badge/license-Apache%202.0-green" alt="Apache 2.0"></a>
<a href="https://huggingface.co/joelbarmettler/md-reheader"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Model-yellow" alt="HuggingFace Model"></a>
<a href="https://huggingface.co/datasets/joelbarmettler/md-reheader-dataset"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Dataset-Explore-yellow" alt="HuggingFace Dataset"></a>
<a href="https://github.com/joelbarmettlerUZH/md-reheader"><img src="https://img.shields.io/github/stars/joelbarmettlerUZH/md-reheader?style=social" alt="GitHub stars"></a>
</p>
</div>
---
## The problem
PDF-to-markdown tools like [MinerU](https://github.com/opendatalab/MinerU), [Docling](https://github.com/DS4SD/docling), and [Marker](https://github.com/VikParuchuri/marker) do great text extraction — then collapse your document structure. Every heading becomes `#` or `##`. TOCs break. RAG chunking breaks. Navigation breaks.
**md-reheader** fixes it. A 0.6B-parameter Qwen3 fine-tune reads the document and predicts the correct H1–H6 level for every heading in a single forward pass.
<p align="center">
<img src="https://raw.githubusercontent.com/joelbarmettlerUZH/md-reheader/main/docs/hero.png" alt="Source PDF → md-reheader → hierarchy restored vs. flat PDF-parser output" width="800">
</p>
---
## Quick start
### CLI
```bash
pip install md-reheader
rehead --input flat.md --output fixed.md
```
Auto-detects CUDA. Use `--cpu` or `--gpu` to override. Omit `--output` to stream to stdout (pipe-friendly).
```bash
rehead -i flat.md | tee fixed.md # pipe
rehead -i flat.md --gpu -o out/fixed.md # creates nested dirs
rehead --help # all flags
```
### Python API
```python
from md_reheader.inference.predict import load_model, reheader_document
model, tokenizer = load_model("joelbarmettler/md-reheader")
flat = open("document.md").read()
fixed = reheader_document(flat, model, tokenizer)
```
The package handles preprocessing (flattening + body stripping) and postprocessing (applying predicted levels back to the original document) automatically.
### Direct `transformers` usage
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("joelbarmettler/md-reheader")
model = AutoModelForCausalLM.from_pretrained(
"joelbarmettler/md-reheader",
dtype=torch.bfloat16,
device_map="auto",
)
messages = [
{"role": "system", "content": "You are a markdown document structure expert. Given a markdown document with incorrect or flattened heading levels, output each heading with its correct markdown prefix (# for level 1, ## for level 2, etc.), one per line."},
{"role": "user", "content": "# Introduction\n\nSome text...\n\n# Background\n\nMore text...\n\n# Methods"},
]
input_text = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True, enable_thinking=False,
)
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(**inputs, max_new_tokens=4096, do_sample=False)
generated = outputs[0][inputs["input_ids"].shape[1]:]
print(tokenizer.decode(generated, skip_special_tokens=True))
# # Introduction
# ## Background
# ## Methods
```
> **Important:** pass `enable_thinking=False` to `apply_chat_template`. Without it, the model enters a repetition loop because training used the non-thinking chat format.
---
### Self-host with vLLM
md-reheader exposes the standard OpenAI-compatible chat endpoint when served with [vLLM](https://github.com/vllm-project/vllm) — higher throughput than raw `transformers`, and drop-in client compatibility.
```bash
pip install vllm
vllm serve joelbarmettler/md-reheader --dtype bfloat16 --max-model-len 8192
```
On <10 GB cards (e.g. RTX 2000/3060), add `--enforce-eager --gpu-memory-utilization 0.70` to skip CUDA-graph allocations that otherwise OOM.
### Remote inference (vLLM or any OpenAI-compatible endpoint)
Once a server is running, use md-reheader as a thin client — no local weights needed.
**CLI:**
```bash
rehead -i flat.md -o fixed.md --endpoint http://localhost:8000/v1
# With auth:
rehead -i flat.md -o fixed.md --endpoint https://api.example.com/v1 --api-key sk-xxx
# or set MD_REHEADER_API_KEY in the environment
```
**Python:**
```python
from md_reheader.inference.remote import reheader_document_remote
fixed = reheader_document_remote(
open("flat.md").read(),
endpoint="http://localhost:8000/v1",
model="joelbarmettler/md-reheader",
api_key=None, # or a bearer token
)
```
The remote client preprocesses locally (flatten + strip), sends a chat completion to the server with `chat_template_kwargs={"enable_thinking": false}` to match training, and applies predicted levels back to the original document. Identical output to local inference.
## How it works
```
flat markdown ──► flatten headings to # ──► strip body to 128+128 tokens
│
▼
restored markdown ◄── apply predicted levels ◄── Qwen3-0.6B (fine-tuned)
```
1. Extract headings with [markdown-it-py](https://github.com/executablebooks/markdown-it-py) (correctly skips code blocks).
2. Flatten every heading to `# ` — the model ignores input levels.
3. Strip each section's body to its first 128 + last 128 tokens — preserves structural cues, kills the context bloat.
4. Qwen3-0.6B predicts the correct `#` prefix per heading.
5. Levels get mapped back to the original document.
---
## Evaluation
Benchmarked on 7,321 held-out documents from GitHub markdown and Wikipedia.
| Metric | All-H1 baseline | Heuristic | **md-reheader** |
|-------------------------|:---------------:|:---------:|:---------------:|
| Exact match | 0.0% | 14.5% | **56.1%** |
| Per-heading accuracy | 13.1% | 49.1% | **80.6%** |
| Hierarchy preservation | 61.3% | 68.6% | **91.0%** |
| Mean absolute error | 1.38 | 0.62 | **0.22** |
### Per-level accuracy
| | H1 | H2 | H3 | H4 | H5 | H6 |
|----------|:---:|:---:|:---:|:---:|:---:|:---:|
| Accuracy | 77% | 85% | 78% | 68% | 45% | 50% |
H1–H3 land in the 77–85% band; H5/H6 drop but still beat baselines. Most deep-level errors are off-by-one — the relative structure survives.
### By document depth
| Max depth | Exact match | Per-heading accuracy | Hierarchy |
|-----------|:-----------:|:--------------------:|:---------:|
| Depth 2 | 83% | 91% | 95% |
| Depth 3 | 54% | 82% | 90% |
| Depth 4 | 32% | 70% | 88% |
| Depth 5-6 | 33% | 65% | 89% |
### By source
| Source | Exact match | Per-heading accuracy |
|-------------------|:-----------:|:--------------------:|
| GitHub markdown | 49.5% | 74.0% |
| Wikipedia | 71.3% | 95.5% |
---
## Speed
| Document size | RTX 4090 (BF16) | CPU (fp32) |
|---------------|:---------------:|:----------:|
| < 1k tokens | 0.4s | 5s |
| 1k–2k tokens | 0.8s | 10s |
| 2k–4k tokens | 1.4s | ~20s |
| 4k–8k tokens | 3.4s | ~60s |
Documents longer than ~8k tokens (after stripping) are truncated from the tail.
---
## Training
| Item | Value |
|-------------------|------------------------------------------------------------|
| Base model | [Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) (text-only) |
| Framework | [Axolotl](https://github.com/axolotl-ai-cloud/axolotl) |
| Training data | ~197k markdown docs (GitHub + Wikipedia, depth 4+ oversampled 2–8×) |
| Hardware | 2× RTX 4090, DDP, BF16 |
| Sequence length | 8,192 tokens with sample packing |
| Learning rate | 5e-5, cosine schedule |
| Epochs | 2 (epoch-1 checkpoint — epoch 2 overfits) |
| Effective batch | 24 |
### Input format during training
1. All headings flattened to `# `.
2. Body text per section truncated to first 128 + last 128 tokens.
3. Document truncated to 8k tokens.
4. Assistant output: one heading per line with its correct `#` prefix.
---
## Limitations
- **Deep nesting (H5/H6):** accuracy drops to 45–50%. Relative structure is preserved but absolute depth gets compressed by 1–2 levels.
- **Ambiguous structure:** heading levels are subjective. The model learns common conventions; it can't resolve genuine ambiguity.
- **Long documents:** >8k tokens (after stripping) get truncated. Headings past the cutoff retain their input levels.
- **English-centric:** trained primarily on English content.
---
## Author
Built by [Joel Barmettler](https://joelbarmettler.xyz/).
## Citation
```bibtex
@software{barmettler2026mdreheader,
author = {Barmettler, Joel},
title = {md-reheader: Restoring Heading Hierarchy in Markdown Documents},
year = {2026},
url = {https://github.com/joelbarmettlerUZH/md-reheader}
}
```

93
chat_template.jinja Normal file
View File

@@ -0,0 +1,93 @@
{%- if tools %}
{{- '<|im_start|>system\n' }}
{%- if messages[0].role == 'system' %}
{{- messages[0].content + '\n\n' }}
{%- endif %}
{{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
{%- for tool in tools %}
{{- "\n" }}
{{- tool | tojson }}
{%- endfor %}
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
{%- else %}
{%- if messages[0].role == 'system' %}
{{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
{#- Determine the real last index: use provided value or default to messages length - 1 #}
{%- if real_last_index is defined and real_last_index is not none %}
{%- set ns.real_last_index = real_last_index %}
{%- else %}
{%- set ns.real_last_index = messages|length - 1 %}
{%- endif %}
{%- for message in messages[::-1] %}
{%- set index = (messages|length - 1) - loop.index0 %}
{%- if ns.multi_step_tool and message.role == "user" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
{%- set ns.multi_step_tool = false %}
{%- set ns.last_query_index = index %}
{%- endif %}
{%- endfor %}
{%- for message in messages %}
{%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
{%- elif message.role == "assistant" %}
{%- set content = message.content %}
{%- set reasoning_content = '' %}
{%- if message.reasoning_content is defined and message.reasoning_content is not none %}
{%- set reasoning_content = message.reasoning_content %}
{%- else %}
{%- if '</think>' in message.content %}
{%- set content = message.content.split('</think>')[-1].lstrip('\n') %}
{%- set reasoning_content = message.content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
{%- endif %}
{%- endif %}
{%- if loop.index0 > ns.last_query_index %}
{%- if loop.index0 == ns.real_last_index or (loop.index0 != ns.real_last_index and reasoning_content) %}
{{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
{%- else %}
{{- '<|im_start|>' + message.role + '\n' + content }}
{%- endif %}
{%- else %}
{{- '<|im_start|>' + message.role + '\n' + content }}
{%- endif %}
{%- if message.tool_calls %}
{%- for tool_call in message.tool_calls %}
{%- if (loop.first and content) or (not loop.first) %}
{{- '\n' }}
{%- endif %}
{%- if tool_call.function %}
{%- set tool_call = tool_call.function %}
{%- endif %}
{{- '<tool_call>\n{"name": "' }}
{{- tool_call.name }}
{{- '", "arguments": ' }}
{%- if tool_call.arguments is string %}
{{- tool_call.arguments }}
{%- else %}
{{- tool_call.arguments | tojson }}
{%- endif %}
{{- '}\n</tool_call>' }}
{%- endfor %}
{%- endif %}
{{- '<|im_end|>\n' }}
{%- elif message.role == "tool" %}
{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
{{- '<|im_start|>user' }}
{%- endif %}
{{- '\n<tool_response>\n' }}
{{- message.content }}
{{- '\n</tool_response>' }}
{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
{{- '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
{{- '<|im_start|>assistant\n' }}
{%- if enable_thinking is defined and enable_thinking is false %}
{{- '<think>\n\n</think>\n\n' }}
{%- else %}
{{- '<think>\n\n' }}
{%- endif %}
{%- endif %}

64
config.json Normal file
View File

@@ -0,0 +1,64 @@
{
"architectures": [
"Qwen3ForCausalLM"
],
"attention_bias": false,
"attention_dropout": 0.0,
"bos_token_id": null,
"dtype": "bfloat16",
"eos_token_id": 151645,
"head_dim": 128,
"hidden_act": "silu",
"hidden_size": 1024,
"initializer_range": 0.02,
"intermediate_size": 3072,
"layer_types": [
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention"
],
"max_position_embeddings": 40960,
"max_window_layers": 28,
"model_type": "qwen3",
"num_attention_heads": 16,
"num_hidden_layers": 28,
"num_key_value_heads": 8,
"pad_token_id": 151643,
"rms_norm_eps": 1e-06,
"sliding_window": null,
"tie_word_embeddings": true,
"transformers_version": "5.3.0",
"use_cache": false,
"use_sliding_window": false,
"vocab_size": 151936,
"rope_theta": 1000000,
"rope_parameters": {
"rope_theta": 1000000,
"rope_type": "default"
}
}

12
generation_config.json Normal file
View File

@@ -0,0 +1,12 @@
{
"do_sample": true,
"eos_token_id": [
151645,
151643
],
"pad_token_id": 151643,
"temperature": 0.6,
"top_k": 20,
"top_p": 0.95,
"transformers_version": "5.3.0"
}

3
model.safetensors Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:c5e504a40a9f69359dcff5c0e4b10433d4bfb943e17ed9dc8d929e0793c2a5a4
size 1192135096

3
tokenizer.json Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
size 11422650

14
tokenizer_config.json Normal file
View File

@@ -0,0 +1,14 @@
{
"add_prefix_space": false,
"backend": "tokenizers",
"bos_token": null,
"clean_up_tokenization_spaces": false,
"eos_token": "<|im_end|>",
"errors": "replace",
"is_local": false,
"model_max_length": 131072,
"pad_token": "<|endoftext|>",
"split_special_tokens": false,
"tokenizer_class": "Qwen2Tokenizer",
"unk_token": null
}