Spaces:
Sleeping
Sleeping
import gradio as gr | |
import torch | |
from CGPT_utils import * | |
from transformers import RobertaTokenizerFast | |
config = { | |
"batch_size" : 8, | |
"num_epochs" : 5, | |
"lr": 10**-4, | |
"seq_len": 80, | |
"d_model" : 128, | |
"src_lang" : "ChEMBL_ID", | |
"tgt_format" : "SMILES", | |
"model_folder" : "weights", | |
"model_basename": "tdmodel_", | |
"preload" : None, | |
"tokenizer_file" : "tokenizer_{0}.json", | |
"experiment_name": "runs/tmodel", | |
"SMILES dataset" : './data/train_dataset.csv', | |
"validation dataset" : './data/test_dataset.csv', | |
"decoder only" : True, | |
} | |
def causal_mask(size): | |
mask = torch.triu(torch.ones((1, size, size)), diagonal=1).type(torch.int) | |
return mask == 0 | |
def greedy_decode(model, source, source_mask, tokenizer_tgt, max_len): | |
sos_idx = tokenizer_tgt.encode("<s>", add_special_tokens=False)[0] | |
# print(sos_idx) | |
eos_idx = tokenizer_tgt.encode("</s>", add_special_tokens=False)[0] | |
# Initialize the decoder input with the sos token | |
decoder_input = torch.cat([torch.empty(1, 1).fill_(sos_idx).type_as(source), source], dim=1) | |
while True: | |
if decoder_input.size(1) == max_len: | |
break | |
# build mask for target | |
decoder_mask = causal_mask(decoder_input.size(1)).type_as(source_mask) | |
# calculate output | |
out = model.decode(decoder_input, decoder_mask) | |
# get next token | |
prob = model.project(out[:, -1]) | |
_, next_word = torch.max(prob, dim=1) | |
decoder_input = torch.cat( | |
[decoder_input, torch.empty(1, 1).type_as(source).fill_(next_word.item())], dim=1 | |
) | |
if next_word == eos_idx: | |
break | |
return decoder_input.squeeze(0) | |
# Specify the directory where the tokenizers are saved | |
chem_tokenizer_dir = "chem_tokenizer" | |
chem_tokenizer = RobertaTokenizerFast.from_pretrained(chem_tokenizer_dir) | |
# Load the model | |
model_checkpoint = torch.load("tdmodel_04.pt", map_location=torch.device('cpu')) | |
model_state_dict = model_checkpoint['model_state_dict'] | |
# Initialize the model | |
model = build_decoder_only_transformer(50265, config['seq_len'], config['d_model']) | |
# Load the model state | |
model.load_state_dict(model_state_dict) | |
def inference(input_sequence): | |
input_tokens = chem_tokenizer.encode(input_sequence, add_special_tokens=False) | |
max_len = 64 | |
input_tensor = torch.tensor(input_tokens).unsqueeze(0) # Add batch dimension | |
source_mask = causal_mask(input_tensor.size(1)).unsqueeze(0) # Assuming decoder mask creation is available | |
# Query the Model | |
output_tensor = greedy_decode(model, input_tensor, source_mask, chem_tokenizer, max_len) | |
# Decode Output | |
output_sequence = chem_tokenizer.decode(output_tensor.cpu().numpy()) | |
# print("Input Sequence:", input_sequence) | |
# print("Model Output:", output_sequence) | |
output_sequence = output_sequence.replace("<s>", "").replace("</s>", "") | |
return output_sequence | |
textinput = gr.components.Textbox(lines=1, label="Enter a SMILES", placeholder="C1=CC=C(C=C1)C") | |
textoutput = gr.components.Textbox() | |
examples = ["C1=CC=C(C=C1)C", "C1CC1C(=O)NC2=CC=CC(=C2)N", "CC(=O)OC1=CC=C"] | |
intf = gr.Interface(fn=inference, inputs=textinput, outputs=textoutput, examples=examples, title="CGPT (Chemical Generative Pretrained Transformer)", description="This model is a decoder-only transformer trained on the entire pubmed dataset. It can be used to complete the SMILES string from a given input SMILES string. d_model is 128, sequence length is 80, 8 attention heads per block with 6 blocks for a total of 16,473,203 parameters") | |
intf.launch(inline=False) |