--- license: gemma library_name: transformers pipeline_tag: text-generation base_model: google/gemma-4-E4B-it datasets: - lamm-mit/graph-canvas-inpainting-121k language: - en tags: - gemma - graph-completion - graph-inpainting - scientific-reasoning - structured-generation - grpo - vllm --- # 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 - Model: - Dataset: - Code: - Base model: - Detailed training and evaluation guide: [`GRAPH_COMPLETION_GRPO_README.md`](https://github.com/lamm-mit/graph-preflexor-grpo/blob/main/GRAPH_COMPLETION_GRPO_README.md) ## 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. ![image](https://cdn-uploads.huggingface.co/production/uploads/623ce1c6b66fedf374859fe7/MxCwVjW4v6kILliJv47Xm.png) ## Output contract The model may use Gemma's native thinking channel. The scored final response is the last complete answer block: ```text { "nodes": [...], "edges": [...] } ``` 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 ```bash 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. ```bash 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: ```bash 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. ```bash huggingface-cli login ``` Verify the relevant interfaces without loading the model: ```bash 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: ```bash 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. ```bash 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. ```bash 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. ```bash 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. ```bash 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. ```bash 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. ```bash 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`: ```json { "nodes": [ {"id": "Humidity"}, {"id": "WaterUptake"}, {"id": "Stiffness"} ], "edges": [ { "source": "Humidity", "relation": "increases", "target": "WaterUptake" } ] } ``` Then replace the inline argument with: ```text --partial_graph_file inputs/humidity-partial-graph.json ``` ## Official-test benchmark Sampling and analysis are deliberately separate: ```text 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 ```bash 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: ```bash 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: ```bash 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: ```bash 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: ```text 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: ```bash 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: ```text 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: ```text 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: ```text 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 ```