From 41ed79e963224bd83711cae1f00e7de15716605c Mon Sep 17 00:00:00 2001 From: ModelHub XC Date: Sat, 13 Jun 2026 10:17:16 +0800 Subject: [PATCH] =?UTF-8?q?=E5=88=9D=E5=A7=8B=E5=8C=96=E9=A1=B9=E7=9B=AE?= =?UTF-8?q?=EF=BC=8C=E7=94=B1ModelHub=20XC=E7=A4=BE=E5=8C=BA=E6=8F=90?= =?UTF-8?q?=E4=BE=9B=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Model: joelbarmettler/md-reheader Source: Original Platform --- .gitattributes | 36 ++++++ README.md | 269 +++++++++++++++++++++++++++++++++++++++++ chat_template.jinja | 93 ++++++++++++++ config.json | 64 ++++++++++ generation_config.json | 12 ++ model.safetensors | 3 + tokenizer.json | 3 + tokenizer_config.json | 14 +++ 8 files changed, 494 insertions(+) create mode 100644 .gitattributes create mode 100644 README.md create mode 100644 chat_template.jinja create mode 100644 config.json create mode 100644 generation_config.json create mode 100644 model.safetensors create mode 100644 tokenizer.json create mode 100644 tokenizer_config.json diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..52373fe --- /dev/null +++ b/.gitattributes @@ -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 diff --git a/README.md b/README.md new file mode 100644 index 0000000..09829aa --- /dev/null +++ b/README.md @@ -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 +--- + +
+ +

md-reheader

+ +

Restore heading hierarchy in markdown documents with a fine-tuned 0.6B LLM.

+ +

+ PyPI + Python 3.12+ + Apache 2.0 + HuggingFace Model + HuggingFace Dataset + GitHub stars +

+ +
+ +--- + +## 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. + +

+ Source PDF → md-reheader → hierarchy restored vs. flat PDF-parser output +

+ +--- + +## 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} +} +``` diff --git a/chat_template.jinja b/chat_template.jinja new file mode 100644 index 0000000..77ea906 --- /dev/null +++ b/chat_template.jinja @@ -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 XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|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('') and message.content.endswith('')) %} + {%- 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 '' in message.content %} + {%- set content = message.content.split('')[-1].lstrip('\n') %} + {%- set reasoning_content = message.content.split('')[0].rstrip('\n').split('')[-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\n' + reasoning_content.strip('\n') + '\n\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 %} + {{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- message.content }} + {{- '\n' }} + {%- 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 %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n\n' }} + {%- endif %} +{%- endif %} diff --git a/config.json b/config.json new file mode 100644 index 0000000..0be8bb0 --- /dev/null +++ b/config.json @@ -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" + } +} \ No newline at end of file diff --git a/generation_config.json b/generation_config.json new file mode 100644 index 0000000..9bda08a --- /dev/null +++ b/generation_config.json @@ -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" +} diff --git a/model.safetensors b/model.safetensors new file mode 100644 index 0000000..83e21f2 --- /dev/null +++ b/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c5e504a40a9f69359dcff5c0e4b10433d4bfb943e17ed9dc8d929e0793c2a5a4 +size 1192135096 diff --git a/tokenizer.json b/tokenizer.json new file mode 100644 index 0000000..c7afbed --- /dev/null +++ b/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506 +size 11422650 diff --git a/tokenizer_config.json b/tokenizer_config.json new file mode 100644 index 0000000..1c1a1e9 --- /dev/null +++ b/tokenizer_config.json @@ -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 +} \ No newline at end of file