entropy/roberta_zinc_decoder

1---2tags: 3- chemistry4- molecule5- drug6---7 8# Roberta Zinc Decoder9 10This model is a GPT2 decoder model designed to reconstruct SMILES strings from embeddings created by the 11[roberta_zinc_480m](https://huggingface.co/entropy/roberta_zinc_480m) model. The decoder model was 12trained on 30m compounds from the [ZINC Database](https://zinc.docking.org/).13 14The decoder model conditions generation on mean pooled embeddings from the encoder model. Mean pooled 15embeddings are used to allow for integration with vector databases, which require fixed length embeddings.16 17Condition embeddings are passed to the decoder model using the `encoder_hidden_states` attribute. 18The standard `GPT2LMHeadModel` does not support generation with encoder hidden states, so this repo 19includes a custom `ConditionalGPT2LMHeadModel`. See example below for how to instantiate the model.20 21```python22import torch23from transformers import AutoModelForCausalLM, RobertaTokenizerFast, RobertaForMaskedLM, DataCollatorWithPadding24 25tokenizer = RobertaTokenizerFast.from_pretrained("entropy/roberta_zinc_480m", max_len=256)26collator = DataCollatorWithPadding(tokenizer, padding=True, return_tensors='pt')27 28encoder_model = RobertaForMaskedLM.from_pretrained('entropy/roberta_zinc_480m')29encoder_model.eval();30 31commit_hash = '0ba58478f467056fe33003d7d91644ecede695a7'32decoder_model = AutoModelForCausalLM.from_pretrained("entropy/roberta_zinc_decoder",33                                                     trust_remote_code=True, revision=commit_hash)34decoder_model.eval();35 36 37smiles = ['Brc1cc2c(NCc3ccccc3)ncnc2s1',38 'Brc1cc2c(NCc3ccccn3)ncnc2s1',39 'Brc1cc2c(NCc3cccs3)ncnc2s1',40 'Brc1cc2c(NCc3ccncc3)ncnc2s1',41 'Brc1cc2c(Nc3ccccc3)ncnc2s1']42 43inputs = collator(tokenizer(smiles))44outputs = encoder_model(**inputs, output_hidden_states=True)45full_embeddings = outputs[1][-1]46mask = inputs['attention_mask']47mean_embeddings = ((full_embeddings * mask.unsqueeze(-1)).sum(1) / mask.sum(-1).unsqueeze(-1))48 49decoder_inputs = torch.tensor([[tokenizer.bos_token_id] for i in range(len(smiles))])50 51hidden_states = mean_embeddings[:,None] # hidden states shape (bs, 1, -1)52 53gen = decoder_model.generate(54              decoder_inputs,55              encoder_hidden_states=hidden_states,56              do_sample=False, # greedy decoding is recommended57              max_length=100, 58              temperature=1.,59              early_stopping=True,60              pad_token_id=tokenizer.pad_token_id,61                         )62 63reconstructed_smiles = tokenizer.batch_decode(gen, skip_special_tokens=True)64```65 66## Model Performance67 68The decoder model was evaluated on a test set of 1m compounds from ZINC. Compounds 69were encoded with the [roberta_zinc_480m](https://huggingface.co/entropy/roberta_zinc_480m) model 70and reconstructed with the decoder model.71 72The following metrics are computed:73* `exact_match` - percent of inputs exactly reconstructed74* `token_accuracy` - percent of output tokens exactly matching input tokens (excluding padding)75* `valid_structure` - percent of generated outputs that resolved to a valid SMILES string76* `tanimoto` - tanimoto similarity between inputs and generated outputs. Excludes invalid structures77* `cos_sim` - cosine similarity between input encoder embeddings and output encoder embeddings78 79`eval_type=full` reports metrics for the full 1m compound test set.80 81`eval_type=failed` subsets metrics for generated outputs that failed to exactly replicate the inputs.82 83 84|eval_type|exact_match|token_accuracy|valid_structure|tanimoto|cos_sim |85|---------|-----------|--------------|---------------|--------|--------|86|full     |0.948277   |0.990704      |0.994278       |0.987698|0.998224|87|failed   |0.000000   |0.820293      |0.889372       |0.734097|0.965668|88 89 90---91license: mit92---93