Ionic Conductivity Analysis

.github/workflows/CI.yaml

@@ -14,11 +14,9 @@ jobs:
           enable-cache: true
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
-      - name: Install deploy key for BattModels/smirk
-        uses: webfactory/ssh-agent@dc588b651fe13675774614f8e6a936a468676387 #v0.9.0
-        with:
-          ssh-private-key: ${{ secrets.SMIRK_PRIVATE_KEY }}
       - name: install project
         run: uv sync --all-extras --dev
       - name: pytest
-        run: uv run pytest
+        env:
+          WANDB_MODE: disabled
+        run: uv run pytest --durations=0

.gitignore

@@ -28,17 +28,25 @@ mist
 Manifest.toml
 LocalPreferences.toml
 .CondaPkg/
-*.jdl2
 
 # Spack
 .spack-env/
 spack.lock
 
 # Output files
 *.json
+*.tar.xz
+*.tar.gz
+*.tar
 *.jsonl
+*.jld2
 *.bson
 *.png
 *.pdf
 *.h5
 *.csv
+*.tar.xz
+*.tar.gz
+*.tar
+*.jsonl
+*.jld2

README.md

@@ -3,6 +3,9 @@ Benchmarking RoBERTa model pre-training on molecular datasets.
 
 # Installation
 
+The following provides installation instructions for the top-level package (`electrolyte_fm`), optional add-ons for our
+various additional analysis and downstream applications (See `opt/`) may require additional configuration.
+
 ## Polaris
 
 1. Install [rust](https://www.rust-lang.org/tools/install) and [uv](https://docs.astral.sh/uv/getting-started/installation/)
@@ -21,11 +24,18 @@ uv sync
 ```
 ## Artemis
 
-
 Same as above except:
 1. Skip loading conda (just use uv)
 2. Ensure a module for CUDA@12.2 exists, may need to install with spack (make sure `buildable: True`)
 
+## Apptainer
+
+0. Install or load from a module [Apptainer](https://apptainer.org/)
+1. Build the image `bash container/build.sh`, once build relocate the image `mv /tmp/mist.sif ./mist.sif`
+2. Run training within the image `apptainer run --nv mist.sif python train.py ...`
+
+> See `submit/dgx.j2` or `submit/delta.j2` for a more complete example of using the container
+
 # Submitting Jobs
 
 ```shell
@@ -35,33 +45,12 @@ source ./activate # Activate Environment
 
 See `submit/submit.py --help` for more info
 
-## Building Apptainer Image
-
-```shell
-apptainer build --fakeroot \
-    --build-arg SSH_AUTH_SOCK=$SSH_AUTH_SOCK \
-    mist.sif mist.def
-```
-
-## Hackathon
-
-Create a file `hack.yaml` and include it as an overlay to `submit.py` (i.e. `./submit/submit.py ... --data hack.yaml ...`).
-Put the following in `hack.yaml`:
-```yaml
-queue: debug
-account: GPU_Hack
-nodes: 2
-walltime: 1:0:0
-train:
-  data.path: /grand/gpu_hack/FoundEnergy/realspace_v3_dev
-```
-
 # Development
 
 ## Pre-commit
 
 We use [pre-commit](https://pre-commit.com) to preform various linting checks on the code. To enable:
 
 1. Install poetry (See above)
-2. Run pre-commit: `pre-commit`
-3. Run before committing: `pre-commit install --allow-missing-config`
+2. Run pre-commit: `uv run pre-commit`
+3. Run before committing: `uv run pre-commit install --allow-missing-config`

electrolyte_fm/data_modules/__init__.py

@@ -1,5 +1,5 @@
 from .molnet_dataset import MolNetDataModule
 from .property_prediction_dataset import PropertyPredictionDataModule
-from .tmqm_dataset import tmQMDataModule
 from .roberta_dataset import RobertaDataSet
-from .mixture_dataset import ComponentDataModule
+from .tmqm_dataset import tmQMDataModule
+from .mixture_dataset import ComponentDataModule, ComponentDataModuleFast

electrolyte_fm/data_modules/isotope_dataset.py

@@ -0,0 +1,28 @@
+from .property_prediction_dataset import PropertyPredictionDataModule
+from .molnet_dataset import train_val_test_split
+from datasets import Dataset, load_dataset
+from pathlib import Path
+
+
+class IsotopeDataModule(PropertyPredictionDataModule):
+    def __init__(self, path: str, **kwargs):
+        # Set default smi_column
+        self.path = path
+        assert Path(self.path).is_file()
+        super().__init__(**kwargs)
+
+    def prepare_data(self):
+        # Fetch data from the head node
+        self.dataset
+
+    def _get_dataset(self):
+        # Load the dataset
+        ds: Dataset = load_dataset(
+            "csv",
+            data_files=[self.path],
+            split="train",
+            keep_in_memory=False,
+            save_infos=False,
+        )  # type: ignore
+
+        return train_val_test_split(ds)