Update README.md

README.md

@@ -13,6 +13,7 @@ tags:
 - high-throughput-screening
 - virtual-drug-screening
 - haicu
+- virtual-human-chc
 library_name: pytorch
 language:
 - en
@@ -26,9 +27,7 @@ MolE learns task-independent molecular representations of chemicals via Graph Is
 
 ## Model versions
 
-- **MolE Antimicrobial Prediction** (March 2024): Pretrained representation model + XGBoost classifier trained on antimicrobial screening data (Maier et al., 2018).
-- Repository: <https://github.com/rolayoalarcon/mole_antimicrobial_potential>
-- Hugging Face Hub: `virtual-human-chc/MolE`
+**MolE Antimicrobial Prediction** (March 2024): Pretrained representation model + XGBoost classifier trained on antimicrobial screening data (Maier et al., 2018).
 
 ## Long description
 
@@ -42,16 +41,17 @@ molecular embeddings from SMILES strings.
 ### Input
 
 - **Description:** Table of chemicals with their SMILES representations.
-- **Shape:** `[n, 2]`
-- **Data format:**
-  - `chem_name` *(str)*: Name of the molecule (e.g., *Halicin*).
-  - `smiles` *(str)*: SMILES representation of the molecule (e.g.,`C1=C(SC(=N1)SC2=NN=C(S2)N)[N+](=O)[O-]`).
+- **Input format:**
+  - **Shape:** `[n, 2]`
+  - **Columns:**
+    - `chem_name` *(str)*: Name of the molecule (e.g., *Halicin*).
+    - `smiles` *(str)*: SMILES representation of the molecule (e.g.,`C1=C(SC(=N1)SC2=NN=C(S2)N)[N+](=O)[O-]`).
 - **Example input file:** `input/examples_molecules.tsv`
 
 ### Model
 
 - **Modality:** Probability of microbial growth inhibition.
-- **Scale:** Combination of bacterial strain and compound.
+- **Scale:** Per bacterial strain and chemical compound.
 - **Description:**
   - The model computes antimicrobial predictive probabilities for **40 bacterial strains** contained in [Maier et al., 2018](https://www.nature.com/articles/nature25979), using a two-step process:
     1.  Generate molecular embeddings with a pre-trained **GINet** representation model (`model.pth`, `config.yaml`).
@@ -132,80 +132,4 @@ print(pred)
 
 ## Copyright
 
-Code derived from <https://github.com/rolayoalarcon/MolE> is licensed under the **MIT License**, © 2024 Roberto Olayo Alarcon. Model weights are licensed under **Creative Commons Attribution 4.0 International (CC BY 4.0)**, © 2024 Roberto Olayo Alarcon. Additional code © 2025 Maksim Pavlov, licensed under MIT.
-
-<!-- # MolE - Antimicrobial Prediction
-
-This model uses MolE's pre-trained representation to train XGBoost models to predict the antimicrobial activity of compounds based on their molecular structure. The model was developed by Roberto Olayo Alarcon et al. and more information can be found in the [GitHub repository](https://github.com/rolayoalarcon/MolE) and the [accompanying paper](https://www.nature.com/articles/s41467-025-58804-4).
-
-## Files:
-
-- `model.pth` - the pre-trained representation model's weights
-- `config.yaml` - model configuration
-- `MolE-XGBoost-08.03.2024_14.20.pkl` - pretrained XGBoost model
-
-## Usage
-
-### Inference Example
-
-Below is a minimal example showing how to load and run inference with **MolE** directly from the Hugging Face Hub.
-
-Necessary dependancies:
-
-```bash
-pip install torch pyyaml pandas huggingface-hub mole_package
-```
-
-```python
-import torch
-import yaml
-import pickle
-import pandas as pd
-from huggingface_hub import hf_hub_download
-from mole_package import ginet_concat, mole_antimicrobial_prediction, mole_representation, dataset_representation
-
-class MolE:
-    def __init__(self, device='auto'):
-        repo = "virtual-human-chc/MolE"
-        self.device = "cuda:0" if device == "auto" and torch.cuda.is_available() else "cpu"
-
-        # Download + load
-        cfg = yaml.safe_load(open(hf_hub_download(repo, "config.yaml")))
-        self.model = ginet_concat.GINet(**cfg["model"]).to(self.device)
-        self.model.load_state_dict(torch.load(hf_hub_download(repo, "model.pth"), map_location=self.device))
-        self.xgb = pickle.load(open(hf_hub_download(repo, "MolE-XGBoost-08.03.2024_14.20.pkl"), "rb"))
-
-    def predict_from_smiles(self, smiles_tsv):
-        smiles_df = mole_representation.read_smiles(smiles_tsv, "smiles", "chem_name")
-        emb = dataset_representation.batch_representation(smiles_df, self.model, "smiles", "chem_name", device=self.device)
-        X_input = mole_antimicrobial_prediction.add_strains(
-            emb, "data/01.prepare_training_data/maier_screening_results.tsv.gz"
-        )
-        probs = self.xgb.predict_proba(X_input)[:, 1]
-        return pd.DataFrame(
-            {"antimicrobial_predictive_probability": probs},
-            index=X_input.index
-        )
-```
-### Run inference:
-
-```python
-mole = MolE()
-pred = mole.predict_from_smiles("examples/input/examples_molecules.tsv")
-print(pred)
-```
-
-## Metadata
-
-### Input
-
-The input is a TSV file with two columns: `chem_name` and `smiles`. The column 'chem_name' contains the name of the molecule from PubChem, e.g. Halicin, and the column 'smiles' contains the chemical formula in SMILES format, e.g. `C1=C(SC(=N1)SC2=NN=C(S2)N)[N+](=O)[O-]`. An example input is the file `examples\input\example_molecules.tsv`.
-
-### Output
-
-The output is a TSV file with two columns: `pred_id` and `antimicrobial_predictive_probability`. The column `pred_id` contains a given molecule and a bacteria, e.g. Halicin:Akkermansia muciniphila (NT5021), and the column `antimicrobial_predictive_probability` contains antimicrobial potential (AP) scores for
-molecule prioritization, reflecting the chance of the given molecule having growth inhibition effect on the corresponding bacteria, e.g. 0.021192694. An example output is `examples/output/example_molecules_prediction.tsv`.
-
-## Copyright
- 
-Code derived from https://github.com/rolayoalarcon/MolE is licensed under the MIT license, Copyright (c) 2024 Roberto Olayo Alarcon. The [model weights](https://doi.org/10.5281/zenodo.10803099) are licensed under [Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0/legalcode), Copyright (c) 2024 Roberto Olayo Alarcon. The other code is licensed under the MIT license, Copyright (c) 2025 Maksim Pavlov. -->
+Code derived from <https://github.com/rolayoalarcon/MolE> is licensed under the **MIT License**, © 2024 Roberto Olayo Alarcon. Model weights are licensed under **Creative Commons Attribution 4.0 International (CC BY 4.0)**, © 2024 Roberto Olayo Alarcon. Additional code © 2025 Maksim Pavlov, licensed under MIT.