load only needed columns

app.py

@@ -12,7 +12,21 @@ from pymatgen.entries.computed_entries import (ComputedStructureEntry,
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 # Load only the train split of the dataset
-dataset = load_dataset("LeMaterial/leDataset", token=HF_TOKEN, split="train")
+dataset = load_dataset(
+    "LeMaterial/leDataset",
+    token=HF_TOKEN,
+    split="train",
+    columns=[
+        "lattice_vectors",
+        "species_at_sites",
+        "cartesian_site_positions",
+        "energy",
+        "energy_corrected",
+        "immutable_id",
+        "elements",
+        "functional",
+    ],
+)
 
 # Convert the train split to a pandas DataFrame
 train_df = dataset.to_pandas()