Modelpublic

public-knowledge-project/agentic-jats-annotation-qwen3.5-9b-lora-v4-rl-step25

sourceHugging Faceapache-2.0updated 4mo agoView on Hugging Face

0likes17downloads

Model Card

Agentic JATS Annotation — Qwen3.5-9B LoRA (SFT + RL v4 step 25)

Single composed LoRA adapter (rank 64, alpha 128) that goes directly on top of Qwen/Qwen3.5-9B (loaded as AutoModelForImageTextToText). It is the mathematical composition of:

—a rank-32 SFT LoRA (3 epochs over ~540 gold tool-call trajectories), plus
—a rank-32 DAPO RL LoRA (25 steps of multi-turn reinforcement learning).

The model converts academic papers (DOCX/PDF → markdown via Docling) into JATS XML by emitting one typed <tool_call>{...}</tool_call> per turn (~20 actions covering front-matter, body structure, references). A companion environment owns the JATS tree and applies actions, so tag-balance and well-formedness are structurally guaranteed.

⚠️ Inference requires both the model and the JATS environment. The model alone emits actions; the environment applies them to build XML state and produces the next observation. The minimal inference-only package (env, tool schemas, serializer, runnable example) is at <https://github.com/parthsarin/agentic-jats-annotation-inference>. Full training source (RL, reward, gold-scratchpad generation) is at <https://github.com/parthsarin/jats-annotation-via-agentic-scratchpad>.

Quickstart

1. Load model + adapter

python

from peft import PeftModel
from 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:

Convert the following document to JATS XML using the available tools.

<source doc_id='12345-67890-1-CE'>
L0001: Intermittent Left Bundle Branch Block
L0002:
L0003: Mansoor Mozayan, MD, PhD, Marc Mugmon, MD
L0004:
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:

<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_PROMPT
from 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, mark_figure
—Back / refs: add_ref (preferred — folds whole reference into one call), start_ref / ref_field / end_ref (alternative per-field path)
—Meta: skip_lines, unassign_lines, finish

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 next
edit. Use <think>...</think> for brief reasoning (<=200 tokens) and then
emit exactly ONE <tool_call>{...}</tool_call> block. Generation stops at
</tool_call>; the env resumes after parsing.

# Tool-call JSON shape

Each 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 are
INTEGERS (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:str

Body:
  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

1.Base: Qwen/Qwen3.5-9B (9B params, hybrid Gated DeltaNet + Sparse MoE, native 262k ctx)
2.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.
3.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, NSAMPLES=16, MAXTURNS=45.
4.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)

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	4/20	body + xref + inline

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.