Package: astrochem-embedding
FeaturesThe goal of this project is to provide off the shelf language models that work
for studies in astrochemistry; the needs for general molecule discovery/chemistry
are different from astrochemistry, such as the emphasis on transient (e.g. open-shell)
molecules and isotopologues.To support these aspects, we provide here light-weight language models (currently just
a GRU seq2seq model) based off ofSELFIESsyntax and PyTorch. Elements of
this project are designed to strike a balance between research agility and use for
production, and a lot of emphasis is placed on reproducibility using PyTorch Lightning
and a general user interface that doesn’t force the user to know how to develop neural networks.The current highlight of this package is theVICGAE, or variance-invariance-covariance
regularized GRU autoencoder (I guess probablyVICGRUAEis more accurate?). I intend to
write this up in a more detailed form in the near future, but the basic premise is this:
a pair of GRUs form a seq2seq model, whose task is to complete SELFIES strings based off
of randomly masked tokens within the molecule. To improve chemical representation learning,
the VIC regularization uses self-supervision to ensure the token embeddings are chemically
descriptive: we encourage variance (e.g. [CH2] is different from [OH]), invariance (e.g.
isotopic substitution should give more or less the same molecule), and covariance (i.e.
minimizing information sharing between embedding dimensions). While the GRU does the actual
SELFIES reconstruction, the VIC regularization is done at the token embedding level.This has been tested on a few simple comparisons with cosine similarity, comparing isotopic
substitution, element substitution (i.e. C/Si/Ge), and functional group replacement; things
seem to work well for these simple cases.RequirementsThis package requires Python 3.8+, as it uses some decorators only available after 3.7.InstallationThe simplest way to getastrochem_embeddingis through PyPI:$pipinstallastrochem_embeddingIf you’re interested in development, want to train your own model,
or make sure you can take advantage of GPU acceleration, I recommend
usingcondafor your environment specification:$condacreate-nastrochem_embeddingpython=3.8$condaactivateastrochem_embedding$pipinstallpoetry$poetryinstall$condainstall-cpytorchtorchtorchvisioncudatoolkit=11.3UsageThe quickest way to get started is by loading a pre-trained model:>>>fromastrochem_embeddingimportVICGAE>>>importtorch>>>model=VICGAE.from_pretrained()>>>model.embed_smiles("c1ccccc1")will return atorch.Tensor. For now the general interface doesn’t
support batching SMILES just yet, and so to operate on many SMILES
strings would simply require looping:>>>smiles=["c1ccccc1","[C]#N","[13c]1ccccc1"]>>>embeddings=torch.stack([model.embed_smiles(s)forsinsmiles])# optionally convert back to NumPy arrays>>>numpy_embeddings=embeddings.numpy()Project StructureThe project filestructure is laid out as such:├── CITATION.cff
├── codecov.yml
├── CODE_OF_CONDUCT.rst
├── CONTRIBUTING.rst
├── data
│   ├── external
│   ├── interim
│   ├── processed
│   └── raw
├── docs
│   ├── codeofconduct.rst
│   ├── conf.py
│   ├── contributing.rst
│   ├── index.rst
│   ├── license.rst
│   ├── reference.rst
│   ├── requirements.txt
│   └── usage.rst
├── environment.yml
├── models
├── notebooks
│   ├── dev
│   ├── exploratory
│   └── reports
├── noxfile.py
├── poetry.lock
├── pyproject.toml
├── README.rst
├── scripts
│   └── train.py
└── src
   └── astrochem_embedding
      ├── __init__.py
      ├── layers
      │   ├── __init__.py
      │   ├── layers.py
      │   └── tests
      │       ├── __init__.py
      │       └── test_layers.py
      ├── __main__.py
      ├── models
      │   ├── __init__.py
      │   ├── models.py
      │   └── tests
      │       ├── __init__.py
      │       └── test_models.py
      ├── pipeline
      │   ├── data.py
      │   ├── __init__.py
      │   ├── tests
      │   │   ├── __init__.py
      │   │   ├── test_data.py
      │   │   └── test_transforms.py
      │   └── transforms.py
      └── utils.pyA brief summary of what each folder is designed for:datacontains copies of the data used for this project. It is recommended to form a pipeline whereby therawdata is preprocessed, serialized tointerim, and when ready for analysis, placed intoprocessed.modelscontains serialized weights intended for distribution, and/or testing.notebookscontains three subfolders:devis for notebook based development,exploratoryfor data exploration, andreportsfor making figures and visualizations for writeup.scriptscontains files that meant for headless routines, generally those with long compute times such as model training and data cleaning.src/astrochem_embeddingcontains the common code base for this project.Code developmentAll of the code used for this project should be contained insrc/astrochem_embedding,
at least in terms of the high-level functionality (i.e. not scripts), and is intended to be
a standalone Python package.The package is structured to match the abstractions for deep learning, specifically PyTorch,
PyTorch Lightning, and Weights and Biases, by separating parts of data structures and processing
and model/layer development.Some concise tenets for developmentWrite unit tests as you go.Commit changes, and commit frequently. Writesemanticgit commits!Formatting is done withblack; don’t fuss about it 😃For new Python dependencies, usepoetry add <package>.For new environment dependencies, useconda env export -f environment.yml.Notes on best practices, particularly regarding CI/CD, can be found in the extensive
documentation for theHypermodern Python Cookiecutterrepository.LicenseDistributed under the terms of theMIT license,Language models for astrochemistryis free and open source software.IssuesIf you encounter any problems,
pleasefile an issuealong with a detailed description.CreditsThis project was generated from@laserkelvin’s PyTorch Project Cookiecutter,
a fork of@cjolowicz’sHypermodern Python Cookiecuttertemplate.