xpanceo-team/mp-20

git clone https://github.com/microsoft/mattergen.git
cd mattergen
git lfs pull -I data-release/mp-20/ --exclude=""
unzip data-release/mp-20/mp_20.zip -d datasets

import pandas as pd
from tqdm.auto import tqdm
from pymatgen.io.cif import CifParser
from datasets import DatasetDict, Dataset

train_path = "./datasets/mp_20/train.csv"
valid_path = "./datasets/mp_20/val.csv"
test_path = "./datasets/mp_20/test.csv"

train_df = pd.read_csv(train_path)
train_structures = [
    CifParser.from_str(s).parse_structures(primitive=True, on_error="ignore")[0]
    for s in tqdm(train_df["cif"], desc="Parsing CIFs", miniters=20)
]
train_df = train_df.drop(columns=['Unnamed: 0', 'cif'])
train_df["structure"] = [s.to_json() for s in tqdm(train_structures, miniters=20)]

valid_df = pd.read_csv(valid_path)
valid_structures = [
    CifParser.from_str(s).parse_structures(primitive=True, on_error="ignore")[0]
    for s in tqdm(valid_df["cif"], desc="Parsing CIFs", miniters=20)
]
valid_df = valid_df.drop(columns=['Unnamed: 0', 'cif'])
valid_df["structure"] = [s.to_json() for s in tqdm(valid_structures, miniters=20)]

test_df = pd.read_csv(test_path)
test_structures = [
    CifParser.from_str(s).parse_structures(primitive=True, on_error="ignore")[0]
    for s in tqdm(test_df["cif"], desc="Parsing CIFs", miniters=20)
]
test_df = test_df.drop(columns=['Unnamed: 0', 'cif'])
test_df["structure"] = [s.to_json() for s in tqdm(test_structures, miniters=20)]



train_dataset = Dataset.from_pandas(train_df)
valid_dataset = Dataset.from_pandas(valid_df)
test_dataset = Dataset.from_pandas(test_df)

dataset_dict = DatasetDict({
    'train': train_dataset,
    'valid': valid_dataset,
    'test': test_dataset
})
dataset_dict.push_to_hub('xpanceo-team/mp-20', private=True)