Add CMG task: Conditional Material Generation (paper task)

demo/condmatgen_tiny/NatureLM_conditional_v2.json

@@ -0,0 +1,52 @@
+[
+  {"elements": "O Te Tm"},
+  {"elements": "Si Al O"},
+  {"elements": "Fe Ni"},
+  {"elements": "Li Co O"},
+  {"elements": "Na Cl"},
+  {"elements": "Ca Ti O"},
+  {"elements": "Mg O"},
+  {"elements": "Zn S"},
+  {"elements": "Cu Fe S"},
+  {"elements": "Ba Ti O"},
+  {"elements": "Sr Ti O"},
+  {"elements": "Ga As"},
+  {"elements": "In P"},
+  {"elements": "Al N"},
+  {"elements": "Si C"},
+  {"elements": "Fe O"},
+  {"elements": "Mn O"},
+  {"elements": "Cr O"},
+  {"elements": "V O"},
+  {"elements": "Ti O"},
+  {"elements": "Li Mn O"},
+  {"elements": "Li Fe P O"},
+  {"elements": "Na Fe O"},
+  {"elements": "K Al Si O"},
+  {"elements": "Ca Mg Si O"},
+  {"elements": "Zr O"},
+  {"elements": "Hf O"},
+  {"elements": "Nb O"},
+  {"elements": "Mo S"},
+  {"elements": "W Se"},
+  {"elements": "Bi Te"},
+  {"elements": "Sb Te"},
+  {"elements": "Pb Te"},
+  {"elements": "Sn Se"},
+  {"elements": "Ge Te"},
+  {"elements": "Cd Se"},
+  {"elements": "Zn Te"},
+  {"elements": "Cu Zn Sn S"},
+  {"elements": "Ag Bi Se"},
+  {"elements": "La Mn O"},
+  {"elements": "Y Ba Cu O"},
+  {"elements": "Nd Fe B"},
+  {"elements": "Sm Co"},
+  {"elements": "Ce O"},
+  {"elements": "Pr Ni O"},
+  {"elements": "Gd Fe O"},
+  {"elements": "Eu Ti O"},
+  {"elements": "Tb Mn O"},
+  {"elements": "Dy Fe O"},
+  {"elements": "Ho Co O"}
+]

demo/fixture_manifest.csv

@@ -12,3 +12,4 @@ rxn_inversion,demo/datasets/rxn_inversion_sample.csv,50,"~45 train / ~5 test",re
 rxn_replacement,demo/datasets/rxn_inversion_sample.csv,50,"~45 train / ~5 test",repo-local fixture,ready,"Same schema as rxn_inversion (MCQ with 4 options)"
 rxn_naming,demo/datasets/rxn_naming_sample.csv,50,"~45 train / ~5 test",repo-local fixture,ready,"Reaction classification into 10 named categories"
 rxn_truefalse,demo/datasets/rxn_truefalse_sample.csv,50,"~45 train / ~5 test",repo-local fixture,ready,"Binary true/false reaction validity"
+condmatgen,demo/condmatgen_tiny,50,"~45 train / ~5 test",repo-local fixture,ready,"JSON with element lists for conditional material generation"

demo/run_fixture_smoke.py

@@ -43,6 +43,9 @@ def main():
         "rxn_truefalse": datasets_dir / "rxn_truefalse_sample.csv",
     }
 
+    if "condmatgen" in CHEMTASKS:
+        task_configs["condmatgen"] = demo_dir / "condmatgen_tiny"
+
     summary = {}
     for task_name, dataset_path in task_configs.items():
         task_class = CHEMTASKS[task_name]

docs/source/modules.rst

@@ -17,4 +17,5 @@ Modules Reference
    tasks/rxn_naming
    tasks/rxn_truefalse
    tasks/template
-   tasks/smi_permute
+   tasks/smi_permute
+   tasks/condmatgen

docs/source/tasks/condmatgen.rst

@@ -0,0 +1,26 @@
+Conditional Material Generation (CMG)
+======================================
+
+.. currentmodule:: open_r1.tasks.condmatgen.condmatgen
+
+ConditionalMaterialGeneration
+-----------------------------
+
+.. autoclass:: ConditionalMaterialGeneration
+   :members:
+   :show-inheritance:
+
+Task Description
+----------------
+
+Given a set of chemical elements, the model proposes a novel crystalline
+compound (element list and space group number). The model wraps its reasoning
+in ``<think>...</think>`` tags and its answer in ``<answer>...</answer>`` tags.
+
+Reward Functions
+----------------
+
+- **accuracy**: multi-component scoring including SMACT validity, element
+  precision, space group validity, and novelty bonus.
+- **format**: checks presence and ordering of think/answer tags, penalizes
+  short reasoning.

recipes/condmatgen.yaml

@@ -0,0 +1,49 @@
+# Model arguments
+model_revision: main
+torch_dtype: bfloat16
+attn_implementation: flash_attention_2
+bf16: true
+tf32: true
+
+# Chemical Task arguments
+chem_task: condmatgen
+dataset_id_or_path: ${MIST_DATA_DIR}/condmatgen
+rewards:
+- accuracy
+
+# Lora Arguments
+# No LoRA is used here
+
+# Training arguments
+max_steps: 1450
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 8
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+learning_rate: 2.0e-6 # 1.0e-6 as in the deepseek math paper 5-e7 from https://hijkzzz.notion.site/unraveling-rlhf-and-its-variants-engineering-insights#147d9a33ecc9806090f3d5c749d31f05
+lr_scheduler_type: cosine
+warmup_ratio: 0.03
+# GRPO specific parameters
+beta: 0.04 # 0.04 as in the deepseek math paper 0.001 from https://hijkzzz.notion.site/unraveling-rlhf-and-its-variants-engineering-insights#147d9a33ecc9806090f3d5c749d31f05
+max_prompt_length: 600
+max_completion_length: 2048
+num_generations: 4
+use_vllm: true
+vllm_device: "cuda:3"
+vllm_gpu_memory_utilization: 0.8
+vllm_max_model_len: 2048
+
+# Logging arguments
+logging_strategy: steps
+logging_steps: 1
+report_to:
+- wandb
+
+save_strategy: "steps"
+save_steps: 25
+seed: 42
+
+# Hugging Face Hub 
+push_to_hub: false
+  # hub_model_id: llama-3-1-8b-math-orca-qlora-10k-ep1 # if not defined same as output_dir

src/open_r1/tasks/__init__.py

@@ -10,6 +10,11 @@
 from .reactions.smi_permute import PermuteSmiles
 from .smiles_understanding.smiles_hydrogen import SmilesHydrogen
 
+try:
+    from .condmatgen.condmatgen import ConditionalMaterialGeneration
+except ImportError:
+    ConditionalMaterialGeneration = None
+
 # Task keys as specified in the task recipes and documentation
 CHEMTASKS = {
     "rxnpred_with_tags": ForwardReactionWithTags,
@@ -26,3 +31,6 @@
     "rxn_naming": Smiles2Name,
     "rxn_truefalse": ReactionTrueFalse,
 }
+
+if ConditionalMaterialGeneration is not None:
+    CHEMTASKS["condmatgen"] = ConditionalMaterialGeneration

src/open_r1/tasks/condmatgen/comps_used_in_sft.json

@@ -0,0 +1 @@
+[]