public-knowledge-project

agentic-jats-annotation-qwen3.5-9b-lora-v4-rl-step25

README

License: apache-2.0

Quickstart

1. Load model + adapter

python

from peft import PeftModelfrom transformers import AutoModelForImageTextToText, AutoTokenizer base = AutoModelForImageTextToText.from_pretrained(    "Qwen/Qwen3.5-9B",    torch_dtype="bfloat16", trust_remote_code=True, device_map="auto",)model = PeftModel.from_pretrained(    base,    "public-knowledge-project/agentic-jats-annotation-qwen3.5-9b-lora-v4-rl-step25",)tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-9B")

2. Prompt format

Each rollout is a multi-turn chat. The system prompt describes the tool catalog (full text below). The first user message contains the line-numbered source markdown and a reference-id table:

text

Convert the following document to JATS XML using the available tools.<source doc_id='12345-67890-1-CE'>L0001: Intermittent Left Bundle Branch BlockL0002:L0003: Mansoor Mozayan, MD, PhD, Marc Mugmon, MDL0004:L0005: Department of Medicine, MedStar Union Memorial Hospital, Baltimore, MD 21218... (one line per markdown line, prefixed L<4-digit-line> with the original text)</source><rid_table>  cit0001 -> 1  cit0002 -> 2  ... (rid → display-label map for citations; usually built from the back-matter       reference list. If annotating from scratch, you can pass an empty table       and let the model emit refs without xref validation.)</rid_table>Emit one <tool_call>{...}</tool_call> per turn. Use short reasoning.

After each assistant turn, the next user message is a compact status block:

text

<status>  open_elements: <sec depth=1 title='Introduction'>  remaining_unassigned: 32/97  emitted: 8 calls</status><result>ok: <p> spans lines 21..23</result>

3. Generation settings

Stop sequence: </tool_call> — the model is trained to emit exactly one <tool_call>{...}</tool_call> block per turn.
Max generation length per turn: 192–256 tokens.
enable_thinking=False: the SFT teacher data never emitted thinking content; passing enable_thinking=True makes the model ramble in <think>...</think> and waste the token budget. Always pass it explicitly:
python
prompt_str = tok.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=False,)
Temperature: 0.7 for sampling, 0.0 for greedy/eval.

python

inputs = tok(prompt_str, return_tensors="pt").to(model.device)out = model.generate(    **inputs, max_new_tokens=192, temperature=0.7,    stop_strings=["</tool_call>"], tokenizer=tok,)assistant_text = tok.decode(out[0, inputs.input_ids.shape[1]:], skip_special_tokens=False)

4. Multi-turn loop sketch (with the env)

python

from src.env import JatsEnv, SYSTEM_PROMPTfrom src.serialize import calls_to_jats_xml env = JatsEnv(extras={    "reward_spec": {"method": "rule", "ground_truth": gold_jats_xml},  # optional    "extra_info": {        "doc_id": doc_id,        "markdown": md_text,        "rid_table": rid_table,   # {"cit0001": "1", ...}        "system_prompt": SYSTEM_PROMPT,    },    "max_turns": 90,})messages, _ = env.init(prompt=None) for turn in range(env.max_turns):    prompt_str = tok.apply_chat_template(        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False,    )    # ... call model.generate(...) -> assistant_text    messages.append({"role": "assistant", "content": assistant_text})    step = env.step(assistant_text)    if step["done"]:        break    messages.extend(step["observations"]) final_xml = calls_to_jats_xml(env.state.emitted_calls,                              md_lines=md_text.splitlines(),                              rid_table=rid_table)

Tool catalog (summary)

20 typed actions, grouped:

Front matter: set_article_title, add_contrib, add_affiliation, start_abstract / end_abstract, add_keyword
Body: mark_section_start, mark_section_end, mark_paragraph, mark_xref, mark_inline (italic / bold / sup / sub / underline), mark_list_start / mark_list_item / mark_list_end, mark_table,

Full Pydantic schemas in src/tools.py on GitHub.

Without the env (just sampling actions)

If you want to inspect raw outputs without building the full environment, you can feed any reasonable system + user prompt and stop at </tool_call>. The model will emit one tool-call JSON per call. But you'll have to validate, apply, and generate observations yourself — easier to pip install the inference package at https://github.com/parthsarin/agentic-jats-annotation-inference and use the env.

System prompt (verbatim)

This is what the model was trained on. Use it as the system message at inference; the line-numbered source markdown and rid_table go in the first user message (see §2 above).

text

You annotate documents with JATS XML by emitting one tool call per turn.The environment owns the XML tree; you only emit JSON describing the nextedit. Use <think>...</think> for brief reasoning (<=200 tokens) and thenemit exactly ONE <tool_call>{...}</tool_call> block. Generation stops at</tool_call>; the env resumes after parsing.# Tool-call JSON shapeEach call is a SINGLE FLAT JSON object whose discriminator field is "name"and whose other fields are the call's arguments at the SAME nesting level(NOT nested under an "arguments" or "params" key). Line numbers areINTEGERS (e.g. 5), not strings (e.g. NOT "L0005").Three correct examples:  <tool_call>{"name": "set_article_title", "line": 1, "title": "German-Austrian Consensus on Charcot Neuroarthropathy"}</tool_call>  <tool_call>{"name": "mark_section_start", "line": 5, "depth": 1, "title": "Introduction"}</tool_call>  <tool_call>{"name": "mark_xref", "line": 7, "target": "1", "ref_type": "bibr", "rid": "cit0001", "head": "guidelines (", "tail": ")."}</tool_call>Common WRONG shapes the parser rejects:  {"command": "...", ...}                     -- key must be "name"  {"name": "...", "arguments": {...}}         -- args are FLAT, not nested  {"name": "...", "line": "L0001"}            -- "line" is int, not "L..."# Tools (signatures: required fields then [optional])Front-matter:  set_article_title           line:int, title:str  add_contrib                 surname:str, given_names:str, [contrib_type, initials, email, aff_rids:list[str]]  add_affiliation             aff_id:str, text:str  start_abstract / end_abstract   (no args)  add_keyword                 text:strBody:  mark_section_start          line:int, depth:int(1..4), title:str, [sec_type]  mark_section_end            depth:int  mark_paragraph              start_line:int, end_line:int  mark_xref                   line:int, target:str, ref_type:"bibr"|"table"|"fig"|"sec"|"aff"|"fn", rid:str, [head, tail]  mark_inline                 line:int, target:str, tag:"italic"|"bold"|"sup"|"sub"|"underline", [head, tail]  mark_list_start             list_type:"bullet"|"order"|"simple"|"alpha-lower"|"alpha-upper"  mark_list_item              start_line:int, end_line:int  mark_list_end               (no args)  mark_table                  start_line:int, end_line:int, [label, caption]  mark_figure                 line:int, [label, caption, graphic_href]Back / ref-list:  add_ref                     rid:str, label:str, publication_type:"journal"|"book"|"chapter"|"conf-proc"|"thesis"|"webpage"|"other", fields:list[{field:str, value:str}]                              Preferred: emits a whole <ref> in one turn instead of start_ref/ref_field*N/end_ref.                              `field` values: surname|given-names|year|article-title|source|volume|issue|fpage|lpage|pub-id-doi|pub-id-pmid|ext-link-uri|...  start_ref                   rid:str, label:str, [publication_type]    # alternative to add_ref, used with ref_field/end_ref  ref_field                   field:str, value:str    # field: surname|given-names|year|article-title|source|volume|issue|fpage|lpage|pub-id-doi|...  end_ref                     (no args)Meta:  skip_lines                  start_line:int, end_line:int, [reason]  unassign_lines              start_line:int, end_line:int  finish                      (no args)# Semantic rules- depth must equal (parent_depth + 1), or 1 at top level.- mark_xref / mark_inline require the line to already be inside a <p>.- mark_xref's `rid` must reference an entry in the rid_table.- finish() requires zero open sections and zero unassigned non-empty lines.

Lineage

Base: Qwen/Qwen3.5-9B (9B params, hybrid Gated DeltaNet + Sparse MoE, native 262k ctx)
SFT-LoRA (rank 32, alpha 64, all-linear): 3 epochs over ~540 gold trajectories produced by replaying gold JATS XML → tool-call sequences via the project's "XML Compositor" pipeline. Each trajectory replays cleanly through the environment.
RL-LoRA (rank 32, alpha 64) trained sequentially on the SFT-merged base: 25 steps of multi-turn DAPO with eps_clip_low=0.2, eps_clip_high=0.32, no KL loss, overlong shaping, zero-variance group filtering, no advantage std-normalization. Per-turn rollout temperature 1.0, N_SAMPLES=16, MAX_TURNS=45.
This adapter: SFT + RL LoRAs composed into a single rank-64 LoRA. The composition is exact — applying this single LoRA on raw Qwen3.5-9B is mathematically equivalent to applying SFT then RL sequentially.

Training data

Source corpus: PKP (Public Knowledge Project) journal articles paired with JATS XML
~810 paired DOCX → markdown → JATS triples; 80/10/10 train/val/test split
Markdown extracted via Docling (DOCX direct path); a Path-B "DOCX → PDF → markdown" shadow set adds format-robustness for inference time
Reward: F_β=0.5 over (parent_tag, child_tag, depth, text_hash) tuples between the model's emitted tool-call sequence and gold JATS body+front+back, with a small recall floor and anti-hack guards

Offline eval (20 validation rollouts, MAX_TURNS=90, temp=0.7)

Table with columns: Model, mean reward, max reward, positive, best behavior
Model	mean reward	max reward	positive	best behavior
Old SFT (pre-line-anchored)	-0.87	+0.71	3/20	spam `add_affiliation`
3-epoch SFT (line-anchored)	-1.65	+0.70	3/20	body annotation
This (SFT + RL step 25)	-1.59	+0.91

Tool distribution at this checkpoint (448 calls across 30 sampled rollouts): mark_paragraph 99 / mark_xref 79 / skip_lines 51 / add_keyword 42 / mark_section_start 39 / add_contrib 36 / add_affiliation 25 / mark_inline 20. Significantly more diverse than SFT-only baselines.

Limitations

Recall is low on long documents — median gold scratchpad has ~100 calls; rollouts truncate at MAX_TURNS and recall drops. Best for documents under ~60 markdown lines.
finish() rarely successful — most rollouts terminate via turn cap or consecutive-error limit. Downstream consumers should serialize partial state via src/serialize.py:calls_to_jats_xml.
Distribution — trained on PKP journal articles; out-of-distribution document types (books, proceedings, theses) likely degrade.

License

Apache-2.0, matching Qwen3.5-9B.

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Model Details

Model Provider

public-knowledge-project

Model Tree

Base

Qwen/Qwen3.5-9B

Adapter

this model

Input Modalities

TextImageVideo

Output Modalities

Text

Supported Functionality

Dedicated Endpoints

Explore FriendliAI today

Get started Talk to an engineer

README

License: apache-2.0

Quickstart

1. Load model + adapter

python

from peft import PeftModelfrom transformers import AutoModelForImageTextToText, AutoTokenizer base = AutoModelForImageTextToText.from_pretrained(    "Qwen/Qwen3.5-9B",    torch_dtype="bfloat16", trust_remote_code=True, device_map="auto",)model = PeftModel.from_pretrained(    base,    "public-knowledge-project/agentic-jats-annotation-qwen3.5-9b-lora-v4-rl-step25",)tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-9B")

2. Prompt format

Each rollout is a multi-turn chat. The system prompt describes the tool catalog (full text below). The first user message contains the line-numbered source markdown and a reference-id table:

text

Convert the following document to JATS XML using the available tools.<source doc_id='12345-67890-1-CE'>L0001: Intermittent Left Bundle Branch BlockL0002:L0003: Mansoor Mozayan, MD, PhD, Marc Mugmon, MDL0004:L0005: Department of Medicine, MedStar Union Memorial Hospital, Baltimore, MD 21218... (one line per markdown line, prefixed L<4-digit-line> with the original text)</source><rid_table>  cit0001 -> 1  cit0002 -> 2  ... (rid → display-label map for citations; usually built from the back-matter       reference list. If annotating from scratch, you can pass an empty table       and let the model emit refs without xref validation.)</rid_table>Emit one <tool_call>{...}</tool_call> per turn. Use short reasoning.

After each assistant turn, the next user message is a compact status block:

text

<status>  open_elements: <sec depth=1 title='Introduction'>  remaining_unassigned: 32/97  emitted: 8 calls</status><result>ok: <p> spans lines 21..23</result>

3. Generation settings

Stop sequence: </tool_call> — the model is trained to emit exactly one <tool_call>{...}</tool_call> block per turn.
Max generation length per turn: 192–256 tokens.
enable_thinking=False: the SFT teacher data never emitted thinking content; passing enable_thinking=True makes the model ramble in <think>...</think> and waste the token budget. Always pass it explicitly:
python
prompt_str = tok.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=False,)
Temperature: 0.7 for sampling, 0.0 for greedy/eval.

python

inputs = tok(prompt_str, return_tensors="pt").to(model.device)out = model.generate(    **inputs, max_new_tokens=192, temperature=0.7,    stop_strings=["</tool_call>"], tokenizer=tok,)assistant_text = tok.decode(out[0, inputs.input_ids.shape[1]:], skip_special_tokens=False)

4. Multi-turn loop sketch (with the env)

python

from src.env import JatsEnv, SYSTEM_PROMPTfrom src.serialize import calls_to_jats_xml env = JatsEnv(extras={    "reward_spec": {"method": "rule", "ground_truth": gold_jats_xml},  # optional    "extra_info": {        "doc_id": doc_id,        "markdown": md_text,        "rid_table": rid_table,   # {"cit0001": "1", ...}        "system_prompt": SYSTEM_PROMPT,    },    "max_turns": 90,})messages, _ = env.init(prompt=None) for turn in range(env.max_turns):    prompt_str = tok.apply_chat_template(        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False,    )    # ... call model.generate(...) -> assistant_text    messages.append({"role": "assistant", "content": assistant_text})    step = env.step(assistant_text)    if step["done"]:        break    messages.extend(step["observations"]) final_xml = calls_to_jats_xml(env.state.emitted_calls,                              md_lines=md_text.splitlines(),                              rid_table=rid_table)

Tool catalog (summary)

20 typed actions, grouped:

Front matter: set_article_title, add_contrib, add_affiliation, start_abstract / end_abstract, add_keyword
Body: mark_section_start, mark_section_end, mark_paragraph, mark_xref, mark_inline (italic / bold / sup / sub / underline), mark_list_start / mark_list_item / mark_list_end, mark_table,

Full Pydantic schemas in src/tools.py on GitHub.

Without the env (just sampling actions)

System prompt (verbatim)

This is what the model was trained on. Use it as the system message at inference; the line-numbered source markdown and rid_table go in the first user message (see §2 above).

text

You annotate documents with JATS XML by emitting one tool call per turn.The environment owns the XML tree; you only emit JSON describing the nextedit. Use <think>...</think> for brief reasoning (<=200 tokens) and thenemit exactly ONE <tool_call>{...}</tool_call> block. Generation stops at</tool_call>; the env resumes after parsing.# Tool-call JSON shapeEach call is a SINGLE FLAT JSON object whose discriminator field is "name"and whose other fields are the call's arguments at the SAME nesting level(NOT nested under an "arguments" or "params" key). Line numbers areINTEGERS (e.g. 5), not strings (e.g. NOT "L0005").Three correct examples:  <tool_call>{"name": "set_article_title", "line": 1, "title": "German-Austrian Consensus on Charcot Neuroarthropathy"}</tool_call>  <tool_call>{"name": "mark_section_start", "line": 5, "depth": 1, "title": "Introduction"}</tool_call>  <tool_call>{"name": "mark_xref", "line": 7, "target": "1", "ref_type": "bibr", "rid": "cit0001", "head": "guidelines (", "tail": ")."}</tool_call>Common WRONG shapes the parser rejects:  {"command": "...", ...}                     -- key must be "name"  {"name": "...", "arguments": {...}}         -- args are FLAT, not nested  {"name": "...", "line": "L0001"}            -- "line" is int, not "L..."# Tools (signatures: required fields then [optional])Front-matter:  set_article_title           line:int, title:str  add_contrib                 surname:str, given_names:str, [contrib_type, initials, email, aff_rids:list[str]]  add_affiliation             aff_id:str, text:str  start_abstract / end_abstract   (no args)  add_keyword                 text:strBody:  mark_section_start          line:int, depth:int(1..4), title:str, [sec_type]  mark_section_end            depth:int  mark_paragraph              start_line:int, end_line:int  mark_xref                   line:int, target:str, ref_type:"bibr"|"table"|"fig"|"sec"|"aff"|"fn", rid:str, [head, tail]  mark_inline                 line:int, target:str, tag:"italic"|"bold"|"sup"|"sub"|"underline", [head, tail]  mark_list_start             list_type:"bullet"|"order"|"simple"|"alpha-lower"|"alpha-upper"  mark_list_item              start_line:int, end_line:int  mark_list_end               (no args)  mark_table                  start_line:int, end_line:int, [label, caption]  mark_figure                 line:int, [label, caption, graphic_href]Back / ref-list:  add_ref                     rid:str, label:str, publication_type:"journal"|"book"|"chapter"|"conf-proc"|"thesis"|"webpage"|"other", fields:list[{field:str, value:str}]                              Preferred: emits a whole <ref> in one turn instead of start_ref/ref_field*N/end_ref.                              `field` values: surname|given-names|year|article-title|source|volume|issue|fpage|lpage|pub-id-doi|pub-id-pmid|ext-link-uri|...  start_ref                   rid:str, label:str, [publication_type]    # alternative to add_ref, used with ref_field/end_ref  ref_field                   field:str, value:str    # field: surname|given-names|year|article-title|source|volume|issue|fpage|lpage|pub-id-doi|...  end_ref                     (no args)Meta:  skip_lines                  start_line:int, end_line:int, [reason]  unassign_lines              start_line:int, end_line:int  finish                      (no args)# Semantic rules- depth must equal (parent_depth + 1), or 1 at top level.- mark_xref / mark_inline require the line to already be inside a <p>.- mark_xref's `rid` must reference an entry in the rid_table.- finish() requires zero open sections and zero unassigned non-empty lines.

Lineage

Base: Qwen/Qwen3.5-9B (9B params, hybrid Gated DeltaNet + Sparse MoE, native 262k ctx)
SFT-LoRA (rank 32, alpha 64, all-linear): 3 epochs over ~540 gold trajectories produced by replaying gold JATS XML → tool-call sequences via the project's "XML Compositor" pipeline. Each trajectory replays cleanly through the environment.
RL-LoRA (rank 32, alpha 64) trained sequentially on the SFT-merged base: 25 steps of multi-turn DAPO with eps_clip_low=0.2, eps_clip_high=0.32, no KL loss, overlong shaping, zero-variance group filtering, no advantage std-normalization. Per-turn rollout temperature 1.0, N_SAMPLES=16, MAX_TURNS=45.
This adapter: SFT + RL LoRAs composed into a single rank-64 LoRA. The composition is exact — applying this single LoRA on raw Qwen3.5-9B is mathematically equivalent to applying SFT then RL sequentially.

Training data

Source corpus: PKP (Public Knowledge Project) journal articles paired with JATS XML
~810 paired DOCX → markdown → JATS triples; 80/10/10 train/val/test split
Markdown extracted via Docling (DOCX direct path); a Path-B "DOCX → PDF → markdown" shadow set adds format-robustness for inference time
Reward: F_β=0.5 over (parent_tag, child_tag, depth, text_hash) tuples between the model's emitted tool-call sequence and gold JATS body+front+back, with a small recall floor and anti-hack guards

Offline eval (20 validation rollouts, MAX_TURNS=90, temp=0.7)

Table with columns: Model, mean reward, max reward, positive, best behavior
Model	mean reward	max reward	positive	best behavior
Old SFT (pre-line-anchored)	-0.87	+0.71	3/20	spam `add_affiliation`
3-epoch SFT (line-anchored)	-1.65	+0.70	3/20	body annotation
This (SFT + RL step 25)	-1.59	+0.91

Limitations

Recall is low on long documents — median gold scratchpad has ~100 calls; rollouts truncate at MAX_TURNS and recall drops. Best for documents under ~60 markdown lines.
finish() rarely successful — most rollouts terminate via turn cap or consecutive-error limit. Downstream consumers should serialize partial state via src/serialize.py:calls_to_jats_xml.
Distribution — trained on PKP journal articles; out-of-distribution document types (books, proceedings, theses) likely degrade.

License

Apache-2.0, matching Qwen3.5-9B.