predict.py

"""
This files includes a predict function for the Tox21.
As an input it takes a list of SMILES and it outputs a nested dictionary with 
SMILES and target names as keys.
"""

#---------------------------------------------------------------------------------------
# Dependencies
from typing import List
import random
from collections import defaultdict
#---------------------------------------------------------------------------------------
class Tox21RandomClassifier():
    """
    A random classifier that assigns a random toxicity score to a given SMILES string.  
    """
    
    def __init__(self):
        
        self.target_names = [
            "NR-AR", 
            "NR-AR-LBD", 
            "NR-AhR", 
            "NR-Aromatase", 
            "NR-ER", 
            "NR-ER-LBD", 
            "NR-PPAR-gamma", 
            "SR-ARE", 
            "SR-ATAD5", 
            "SR-HSE", 
            "SR-MMP", 
            "SR-p53"
        ]

    def predict(self, smiles_list:List[str]) -> dict:
        """
        Predicts all Tox21 targets for a given list of SMILES strings by assigning
        random toxicity scores.
        """

        predictions = defaultdict(dict)
        for smiles in smiles_list:
            for target in self.target_names:
                predictions[smiles][target] = random.random()
        return predictions

def predict(smiles_list: List[str]) -> dict:
    """
    Applies the classifier to a list of SMILES strings.
    """
    model = Tox21RandomClassifier()
    return model.predict(smiles_list)