Graph-PRefLexOR-4B-Inpainting

lamm-mit/Graph-PRefLexOR-4B-Inpainting is trained to complete and repair scientific mechanism graphs. Given a natural language condition and an empty, incomplete, or corrupted graph canvas, the model returns the complete corrected graph as structured JSON.

Links

What the model was trained to do

The training data contains six graph-inpainting and repair modes:

Mode Input canvas Intended operation
prior_empty No nodes or edges Construct a complete graph from the condition
fixed_nodes_only A fixed subset of nodes and no edges Preserve supplied nodes and add the missing graph
missing_edges Complete nodes and an incomplete edge set Add missing edges while preserving supplied content
partial_subgraph A correct but incomplete subgraph Add missing nodes and edges
wrong_relations Correct nodes/endpoints with corrupted relation labels Repair incorrect relations
extra_edges Correct graph plus spurious edges Remove unsupported edges

For additive completion, supplied objects should normally be marked fixed. For repair or removal, the supplied graph must remain editable.

Output contract

The model may use Gemma's native thinking channel. The scored final response is the last complete answer block:

<answer>
{
  "nodes": [...],
  "edges": [...]
}
</answer>

The graph object must contain:

  • nodes: objects with at least an id;
  • edges: objects with source, relation, and target;
  • no duplicate nodes or edges;
  • no edges whose endpoints are absent from nodes.

The model returns the complete graph, not a patch or edit list. Any optional node or edge payload fields supplied in the canvas are part of the object contract and should be preserved.

Checkpoint-selection results

The following values were measured with deterministic decoding on the fixed 512-task held-out validation selection used for checkpoint comparison. Metrics are source-macro means. These are validation results, not official-test results.

Model Shaped reward Valid completion Exact graph Node score Edge score
google/gemma-4-E4B-it 0.3871 0.6436 0.0597 0.5871 0.3516
Graph-PRefLexOR-4B-Inpainting 0.5460 0.9869 0.1694 0.7124 0.4115

Use the reproducible official-test workflow below to generate test results for the installed model and dataset revisions.

Installation

The provided CLI uses vLLM for generation and the repository's deterministic Python evaluator for scoring.

1. Clone the repository

git clone https://github.com/lamm-mit/graph-preflexor-grpo.git
cd graph-preflexor-grpo

2. Create an environment

The following is a practical GPU inference environment. Install the PyTorch build appropriate for the CUDA driver on the target machine.

python3 -m venv .venv
source .venv/bin/activate

python -m pip install --upgrade pip setuptools wheel
python -m pip install --upgrade \
  torch torchvision torchaudio

python -m pip install --upgrade \
  "transformers>=5.10.1" \
  accelerate datasets huggingface_hub safetensors \
  sentencepiece protobuf tqdm numpy pandas matplotlib

python -m pip install --upgrade vllm

For training or adapter manipulation, also install:

python -m pip install --upgrade peft trl wandb

The inference CLI was written for the current vLLM generation interface, including vLLM 0.23-style prompt truncation. If vLLM and Transformers are already installed in a working training environment, reuse that environment.

3. Authenticate

The release is public, but Gemma-derived weights remain subject to the Gemma license and Hugging Face access requirements.

huggingface-cli login

Verify the relevant interfaces without loading the model:

python -c "import torch,transformers,vllm,datasets; print(torch.__version__, transformers.__version__, vllm.__version__, datasets.__version__)"
python src/sample_graph_completion.py --help
python src/analyze_graph_completion_benchmark.py --help

Quick dataset smoke test: all six modes

This selects one deterministic example from every trained corruption mode, prints the raw model continuation, and then prints the reference graph and score breakdown:

mkdir -p outputs/graph_completion

python -u src/sample_graph_completion.py \
  --model lamm-mit/Graph-PRefLexOR-4B-Inpainting \
  --dataset lamm-mit/graph-canvas-inpainting-121k \
  --split test \
  --modes prior_empty,fixed_nodes_only,missing_edges,partial_subgraph,wrong_relations,extra_edges \
  --num_tasks 6 \
  --num_generations 1 \
  --generation_batch_size 6 \
  --temperature 0.0 \
  --top_p 1.0 \
  --seed 42 \
  --max_prompt_length 4096 \
  --max_completion_length 4096 \
  --dtype bfloat16 \
  --vllm_gpu_memory_utilization 0.45 \
  --chat_template_enable_thinking true \
  --view scored \
  --reward_stage shaped \
  --output_jsonl outputs/graph_completion/public-model-six-mode-smoke.jsonl

Everything between RAW DECODED MODEL COMPLETION and END RAW COMPLETION is the untouched generated string. The CLI does not repair, extract, or rewrite it before scoring.

Manual graph completion

Manual inference takes:

  1. a natural-language scientific condition;
  2. an inline graph through --partial_graph_json, or a JSON file through --partial_graph_file;
  3. the corruption mode;
  4. a fixed-object policy.

Use:

  • --manual_fixed_policy all for prior_empty, fixed_nodes_only, missing_edges, and partial_subgraph. Every supplied node and edge is rendered as [FIXED] and must be preserved exactly.
  • --manual_fixed_policy none for wrong_relations and extra_edges, because the model must be allowed to edit or remove supplied edges.

Manual inputs do not have a gold reference, so they support --view raw, not --view scored.

The six examples below use one common humidity–silk-film mechanism so the difference between task modes is explicit.

1. Construct from an empty prior: prior_empty

The condition supplies all semantic context; the graph canvas is empty.

mkdir -p outputs/graph_completion/manual

python -u src/sample_graph_completion.py \
  --model lamm-mit/Graph-PRefLexOR-4B-Inpainting \
  --condition "Construct a mechanism graph explaining how increasing humidity reduces the stiffness of a silk fibroin film through water uptake, plasticization, and increased chain mobility." \
  --partial_graph_json '{"nodes":[],"edges":[]}' \
  --manual_mode prior_empty \
  --manual_fixed_policy all \
  --num_generations 1 \
  --temperature 0.0 \
  --top_p 1.0 \
  --seed 11 \
  --max_prompt_length 4096 \
  --max_completion_length 4096 \
  --dtype bfloat16 \
  --vllm_gpu_memory_utilization 0.45 \
  --chat_template_enable_thinking true \
  --view raw \
  --output_jsonl outputs/graph_completion/manual/prior-empty.jsonl \
  --output_text_file outputs/graph_completion/manual/prior-empty.txt

Expected operation: introduce the scientifically appropriate nodes and edges needed to represent the mechanism.

2. Connect fixed nodes: fixed_nodes_only

All supplied nodes are fixed and the edge set is empty.

python -u src/sample_graph_completion.py \
  --model lamm-mit/Graph-PRefLexOR-4B-Inpainting \
  --condition "Connect the supplied concepts into a mechanism graph explaining how increasing humidity reduces silk fibroin film stiffness through water uptake and chain mobility." \
  --partial_graph_json '{"nodes":[{"id":"Humidity"},{"id":"WaterUptake"},{"id":"ChainMobility"},{"id":"Stiffness"}],"edges":[]}' \
  --manual_mode fixed_nodes_only \
  --manual_fixed_policy all \
  --num_generations 1 \
  --temperature 0.0 \
  --top_p 1.0 \
  --seed 11 \
  --max_prompt_length 4096 \
  --max_completion_length 4096 \
  --dtype bfloat16 \
  --vllm_gpu_memory_utilization 0.45 \
  --chat_template_enable_thinking true \
  --view raw \
  --output_jsonl outputs/graph_completion/manual/fixed-nodes-only.jsonl \
  --output_text_file outputs/graph_completion/manual/fixed-nodes-only.txt

Expected operation: preserve the four named nodes exactly and infer the mechanistic edge structure. Additional scientifically necessary content may be added.

3. Add omitted relations: missing_edges

The full node set and one correct edge are fixed; other mechanistic edges are missing.

python -u src/sample_graph_completion.py \
  --model lamm-mit/Graph-PRefLexOR-4B-Inpainting \
  --condition "Complete the mechanism graph: humidity increases water uptake, water plasticizes the silk matrix and increases chain mobility, and increased mobility reduces stiffness." \
  --partial_graph_json '{"nodes":[{"id":"Humidity"},{"id":"WaterUptake"},{"id":"ChainMobility"},{"id":"Stiffness"}],"edges":[{"source":"Humidity","relation":"increases","target":"WaterUptake"}]}' \
  --manual_mode missing_edges \
  --manual_fixed_policy all \
  --num_generations 1 \
  --temperature 0.0 \
  --top_p 1.0 \
  --seed 11 \
  --max_prompt_length 4096 \
  --max_completion_length 4096 \
  --dtype bfloat16 \
  --vllm_gpu_memory_utilization 0.45 \
  --chat_template_enable_thinking true \
  --view raw \
  --output_jsonl outputs/graph_completion/manual/missing-edges.jsonl \
  --output_text_file outputs/graph_completion/manual/missing-edges.txt

Expected operation: preserve all supplied objects and add the missing relations.

4. Complete a partial subgraph: partial_subgraph

Only part of the mechanism is present. The supplied subgraph is correct and fixed, but missing intermediate concepts and relations may be added.

python -u src/sample_graph_completion.py \
  --model lamm-mit/Graph-PRefLexOR-4B-Inpainting \
  --condition "Complete a mechanism graph explaining how humidity changes the stiffness of a silk fibroin film through water uptake, plasticization, microstructural disruption, and chain mobility." \
  --partial_graph_json '{"nodes":[{"id":"Humidity"},{"id":"WaterUptake"},{"id":"Stiffness"}],"edges":[{"source":"Humidity","relation":"increases","target":"WaterUptake"}]}' \
  --manual_mode partial_subgraph \
  --manual_fixed_policy all \
  --num_generations 1 \
  --temperature 0.0 \
  --top_p 1.0 \
  --seed 11 \
  --max_prompt_length 4096 \
  --max_completion_length 4096 \
  --dtype bfloat16 \
  --vllm_gpu_memory_utilization 0.45 \
  --chat_template_enable_thinking true \
  --view raw \
  --output_jsonl outputs/graph_completion/manual/partial-subgraph.jsonl \
  --output_text_file outputs/graph_completion/manual/partial-subgraph.txt

Expected operation: retain the supplied subgraph and add missing nodes and edges to produce a complete mechanism.

5. Repair corrupted relation labels: wrong_relations

The endpoints are plausible, but the supplied relation directions contradict the condition. The graph is intentionally editable.

python -u src/sample_graph_completion.py \
  --model lamm-mit/Graph-PRefLexOR-4B-Inpainting \
  --condition "Repair the graph so it states that increasing humidity increases water uptake, while greater water uptake reduces the stiffness of the silk fibroin film." \
  --partial_graph_json '{"nodes":[{"id":"Humidity"},{"id":"WaterUptake"},{"id":"Stiffness"}],"edges":[{"source":"Humidity","relation":"decreases","target":"WaterUptake"},{"source":"WaterUptake","relation":"increases","target":"Stiffness"}]}' \
  --manual_mode wrong_relations \
  --manual_fixed_policy none \
  --num_generations 1 \
  --temperature 0.0 \
  --top_p 1.0 \
  --seed 11 \
  --max_prompt_length 4096 \
  --max_completion_length 4096 \
  --dtype bfloat16 \
  --vllm_gpu_memory_utilization 0.45 \
  --chat_template_enable_thinking true \
  --view raw \
  --output_jsonl outputs/graph_completion/manual/wrong-relations.jsonl \
  --output_text_file outputs/graph_completion/manual/wrong-relations.txt

Expected operation: replace the incorrect relation labels while retaining the scientifically relevant nodes and endpoint pairs.

6. Remove a spurious edge: extra_edges

The graph contains a reverse-causal edge from stiffness to humidity that is not supported by the condition. The graph is intentionally editable.

python -u src/sample_graph_completion.py \
  --model lamm-mit/Graph-PRefLexOR-4B-Inpainting \
  --condition "Repair the mechanism graph so it represents humidity-driven water uptake, increased chain mobility, and reduced silk fibroin film stiffness. Remove unsupported causal edges." \
  --partial_graph_json '{"nodes":[{"id":"Humidity"},{"id":"WaterUptake"},{"id":"ChainMobility"},{"id":"Stiffness"}],"edges":[{"source":"Humidity","relation":"increases","target":"WaterUptake"},{"source":"WaterUptake","relation":"increases","target":"ChainMobility"},{"source":"ChainMobility","relation":"decreases","target":"Stiffness"},{"source":"Stiffness","relation":"increases","target":"Humidity"}]}' \
  --manual_mode extra_edges \
  --manual_fixed_policy none \
  --num_generations 1 \
  --temperature 0.0 \
  --top_p 1.0 \
  --seed 11 \
  --max_prompt_length 4096 \
  --max_completion_length 4096 \
  --dtype bfloat16 \
  --vllm_gpu_memory_utilization 0.45 \
  --chat_template_enable_thinking true \
  --view raw \
  --output_jsonl outputs/graph_completion/manual/extra-edges.jsonl \
  --output_text_file outputs/graph_completion/manual/extra-edges.txt

Expected operation: retain supported content and remove the unsupported Stiffness -> Humidity edge.

Supplying larger graphs through a file

For a larger graph, save a JSON object containing nodes and edges:

{
  "nodes": [
    {"id": "Humidity"},
    {"id": "WaterUptake"},
    {"id": "Stiffness"}
  ],
  "edges": [
    {
      "source": "Humidity",
      "relation": "increases",
      "target": "WaterUptake"
    }
  ]
}

Then replace the inline argument with:

--partial_graph_file inputs/humidity-partial-graph.json

Official-test benchmark

Sampling and analysis are deliberately separate:

model + official test tasks
    -> raw prediction JSONL
    -> deterministic reference scoring
    -> metrics, tables, and publication figures

The current filtered official test split contains 3,641 unique tasks across all six modes.

1. Generate predictions

cd graph-preflexor-grpo
source .venv/bin/activate

mkdir -p outputs/graph_completion/public-release-test

python -u src/sample_graph_completion.py \
  --model lamm-mit/Graph-PRefLexOR-4B-Inpainting \
  --dataset lamm-mit/graph-canvas-inpainting-121k \
  --split test \
  --invalid_pair_policy filter \
  --num_tasks 3641 \
  --num_generations 1 \
  --generation_batch_size 64 \
  --stream_output_jsonl \
  --temperature 0.0 \
  --top_p 1.0 \
  --seed 42 \
  --max_prompt_length 4096 \
  --max_completion_length 4096 \
  --dtype bfloat16 \
  --vllm_gpu_memory_utilization 0.45 \
  --chat_template_enable_thinking true \
  --view raw \
  --output_jsonl outputs/graph_completion/public-release-test/predictions.jsonl \
  > outputs/graph_completion/public-release-test/generation.log 2>&1

Monitor progress:

watch -n 10 'wc -l outputs/graph_completion/public-release-test/predictions.jsonl'
tail -f outputs/graph_completion/public-release-test/generation.log

Generation is complete when the JSONL contains 3,641 lines.

If generation was interrupted, preserve completed rows and generate only missing task identities:

python -u src/sample_graph_completion.py \
  --model lamm-mit/Graph-PRefLexOR-4B-Inpainting \
  --dataset lamm-mit/graph-canvas-inpainting-121k \
  --split test \
  --invalid_pair_policy filter \
  --num_tasks 3641 \
  --num_generations 1 \
  --generation_batch_size 64 \
  --stream_output_jsonl \
  --resume_output_jsonl \
  --temperature 0.0 \
  --top_p 1.0 \
  --seed 42 \
  --max_prompt_length 4096 \
  --max_completion_length 4096 \
  --dtype bfloat16 \
  --vllm_gpu_memory_utilization 0.45 \
  --chat_template_enable_thinking true \
  --view raw \
  --output_jsonl outputs/graph_completion/public-release-test/predictions.jsonl \
  >> outputs/graph_completion/public-release-test/generation.log 2>&1

2. Analyze predictions

This stage does not load the model or run inference:

python -u src/analyze_graph_completion_benchmark.py \
  --predictions outputs/graph_completion/public-release-test/predictions.jsonl \
  --dataset lamm-mit/graph-canvas-inpainting-121k \
  --split test \
  --invalid_pair_policy filter \
  --seed 42 \
  --max_completion_length 4096 \
  --reward_stage shaped \
  --aggregation_unit source \
  --bootstrap_samples 2000 \
  --bootstrap_seed 42 \
  --confidence 0.95 \
  --png_dpi 450 \
  --label "Graph-PRefLexOR-4B-Inpainting" \
  --output_dir outputs/graph_completion/public-release-test/analysis

The analysis directory contains:

scored_predictions.jsonl
benchmark_summary.json
end_to_end_metrics.csv
conditional_graph_metrics.csv
reward_components.csv
generation_metrics.csv
README.md
benchmark_overall.{svg,png}
benchmark_by_mode.{svg,png}
benchmark_conditional_precision_recall.{svg,png}
benchmark_reward_components.{svg,png}
benchmark_generation_diagnostics.{svg,png}

Inspect the main overall metrics:

jq '.end_to_end.overall |
  {
    reward,
    valid_completion,
    exact_match,
    fixed_contract,
    node,
    edge,
    mode_primary
  }' \
  outputs/graph_completion/public-release-test/analysis/benchmark_summary.json

The integrity section should report:

expected_unique_tasks: 3641
predicted_unique_tasks: 3641
predictions_scored: 3641
predictions_without_reference: 0
predictions_with_ambiguous_reference: 0
missing_mode_rows: 0
duplicate_task_generation_rows: 0
missing_expected_tasks: 0

How evaluation works

Evaluation is symbolic and deterministic; it does not use embedding similarity or an LLM judge.

  • Node identity is based on exact node IDs and payloads.
  • Edge identity is based on exact (source, relation, target) content and payloads.
  • Node and edge array order does not affect canonical graph matching.
  • exact_match requires the full canonical graph to equal the reference.
  • Fixed objects must be preserved exactly.
  • Duplicate or dangling objects are structural errors.
  • Mode-specific metrics separately measure additions, relation repair, and spurious-edge removal.
  • Source-macro aggregation prevents sources with more corrupted variants from dominating the reported mean.

The default shaped reward combines:

0.10 format and parsing
0.10 schema and structural validity
0.15 exact fixed-object preservation
0.10 node F1
0.15 edge F1
0.15 mode-specific primary score
0.10 improvement over the unchanged input
0.15 exact canonical graph match

Training summary

  • Base model: google/gemma-4-E4B-it
  • Training dataset: lamm-mit/graph-canvas-inpainting-121k
  • Objective: Group Relative Policy Optimization (GRPO)
  • Selected release: checkpoint 1,350
  • Rollouts per prompt: 8
  • Rollout temperature: 0.8
  • Rollout top-p: 0.95
  • Reward: shaped graph-completion reward
  • Loss: DAPO
  • KL coefficient: 0
  • LoRA rank: 64
  • LoRA alpha: 128
  • LoRA dropout: 0
  • Numeric format: BF16
  • Rollout backend: colocated vLLM
  • Training attention: SDPA
  • vLLM CUDA graphs: disabled during training
  • Chat template: native thinking enabled
  • Maximum prompt length: 4,096 tokens
  • Maximum completion length: 4,096 tokens

The public release contains merged full-model weights for direct inference.

Limitations and responsible use

  • A schema-valid graph is not necessarily scientifically correct.
  • The model can omit mechanisms, introduce unsupported concepts, or choose an incorrect causal relation.
  • Exact symbolic metrics do not award semantic equivalence between differently named nodes or relations.
  • Manual inference has no reference graph and therefore cannot automatically establish scientific correctness.
  • The six corruption-mode labels are part of the task specification; an incorrect mode can encourage the wrong edit behavior.
  • Use --manual_fixed_policy all when user-supplied facts must not be changed.
  • Outputs should be reviewed by a domain expert before use in scientific analysis, education, or decision-making.
  • The model is intended for research on structured scientific reasoning and graph completion, not as an autonomous scientific authority.

License

This model is derived from Gemma and is distributed subject to the Gemma terms and the licenses or terms associated with the training data and source code. Users are responsible for reviewing the applicable terms before redistribution or deployment.

Citation

If you use the model, please cite the model, dataset, and repository:

Graph-PRefLexOR-4B-Inpainting
https://huggingface.co/lamm-mit/Graph-PRefLexOR-4B-Inpainting

Graph Canvas Inpainting 121K
https://huggingface.co/datasets/lamm-mit/graph-canvas-inpainting-121k

Graph-PRefLexOR GRPO
https://github.com/lamm-mit/graph-preflexor-grpo
Downloads last month
264
Safetensors
Model size
8B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for lamm-mit/Graph-PRefLexOR-4B-Inpainting

Finetuned
(262)
this model

Dataset used to train lamm-mit/Graph-PRefLexOR-4B-Inpainting