Base model: meta-llama/Llama-3.1-8B-Instruct
Method: Supervised fine-tuning (SFT) followed by reinforcement learning (GRPO)
Task: Given a funding statement, extract a structured JSON array of funders with their award IDs
Training data: cometadata/funding-extraction-sft-data

Training Pipeline

Stage 1: Supervised Fine-Tuning

Data: cometadata/funding-extraction-sft-data — 1,316 real examples (train.jsonl) + 2,531 synthetic examples (synthetic.jsonl) upsampled 2x = 6,378 total
Epochs: 2
LoRA rank: 128
LoRA alpha: 32
Learning rate: ~2.86e-4
Batch size: 128
Max sequence length: 4,096 tokens
LR schedule: Linear decay
Renderer: llama3 (Llama 3.1 Instruct chat template)
Train on: Last assistant message only

Stage 2: Reinforcement Learning (GRPO)

Algorithm: Group Relative Policy Optimization (GRPO)
Starting checkpoint: SFT final weights
Data: 3,462 train / 385 eval examples
Learning rate: 3e-5
Temperature: 0.8
Batch size: 16, Group size: 8
KL penalty: 0.03 (against SFT reference policy)
Best checkpoint: Step 70 / 217 (selected by eval reward)
Eval reward at best step: 0.959

Reward Function

See https://github.com/cometadata/funding-metadata-enrichment/tree/main/train for the full training code

Gated, hierarchical matching on funder using the Hungarian algorithm for limiting subordinate fields 1:1 funder pairing:

Funder name - Fuzzy matching using a token-sort with acronym and containment boosts, F0.5 score, weight 0.50
Award IDs - Normalized exact matching, F0.5 score, weight 0.50
Funding scheme - Not weighted
Award title - Not weighted

Usage

python

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
model = PeftModel.from_pretrained(base_model, "cometadata/funding-parsing-lora-Llama_3.1_8B-instruct-ep2-r128-a32-grpo")
tokenizer = AutoTokenizer.from_pretrained("cometadata/funding-parsing-lora-Llama_3.1_8B-instruct-ep2-r128-a32-grpo")

messages = [
    {"role": "system", "content": "Extract funding information from the text. Return a JSON array of funders."},
    {"role": "user", "content": "Extract funding information from the following statement:\n\nThis work was supported by the National Science Foundation (Grant No. 2045678) and the European Research Council (ERC-2021-StG-101039567)."}
]

inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True)
outputs = model.generate(inputs, max_new_tokens=512, temperature=0.1)
print(tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True))

Expected output:

json

[
  {
    "funder_name": "National Science Foundation",
    "awards": [
      {
        "award_ids": ["2045678"],
        "funding_scheme": [],
        "award_title": []
      }
    ]
  },
  {
    "funder_name": "European Research Council",
    "awards": [
      {
        "award_ids": ["ERC-2021-StG-101039567"],
        "funding_scheme": [],
        "award_title": []
      }
    ]
  }
]

Training Infrastructure

Trained on Tinker by Thinking Machines Lab

Eval Results

Step	Eval Reward	Notes
0	0.945	SFT baseline
20	0.957
30	0.958
50	0.960
70	0.959	Best checkpoint (this release)
110	0.958
217	0.950	Final step