filter library rather than pandas memory load

app.py

@@ -15,7 +15,11 @@ from pymatgen.entries.computed_entries import (
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-subsets = ["compatible_pbe", "compatible_pbesol", "compatible_scan",]
+subsets = [
+    "compatible_pbe",
+    "compatible_pbesol",
+    "compatible_scan",
+]
 
 # Load only the train split of the dataset
 
@@ -39,11 +43,11 @@ for subset in subsets:
     datasets.append(dataset["train"])
 
 # Convert the train split to a pandas DataFrame
-df = pd.concat([x.to_pandas() for x in datasets])
-train_df = dataset.to_pandas()
-del dataset
-
+# df = pd.concat([x.to_pandas() for x in datasets])
+# train_df = dataset.to_pandas()
+# del dataset
 
+dataset = concatenate_datasets(datasets)
 
 
 def create_phase_diagram(
@@ -60,17 +64,23 @@ def create_phase_diagram(
 
     # Filter entries based on functional
     if functional == "PBE":
-        entries_df = train_df[train_df["functional"] == "pbe"]
+        ds_filter = dataset.filter(lambda example: example["functional"] == "pbe")
+        # entries_df = train_df[train_df["functional"] == "pbe"]
     elif functional == "PBESol":
-        entries_df = train_df[train_df["functional"] == "pbesol"]
+        ds_filter = dataset.filter(lambda example: example["functional"] == "pbesol")
+        # entries_df = train_df[train_df["functional"] == "pbesol"]
     elif functional == "SCAN":
-        entries_df = train_df[train_df["functional"] == "scan"]
+        ds_filter = dataset.filter(lambda example: example["functional"] == "scan")
+        # entries_df = train_df[train_df["functional"] == "scan"]
 
     isubset = lambda x: set(x).issubset(element_list)
     isintersection = lambda x: len(set(x).intersection(element_list)) > 0
-    entries_df = entries_df[
-        [isintersection(l) and isubset(l) for l in entries_df.elements.values.tolist()]
-    ]
+    ds_filter = ds_filter.filter(
+        lambda example: isintersection(example["functional"])
+        and isubset(example["functional"])
+    )
+
+    entries_df = ds_filter.to_pandas()
 
     # Fetch all entries from the Materials Project database
     entries = [