Views
No views yet
1import torch
2from transformers import AutoTokenizer, T5ForConditionalGeneration
3
4torch_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
5tokenizer = AutoTokenizer.from_pretrained("laituan245/molt5-large", model_max_length=512)
6model = T5ForConditionalGeneration.from_pretrained("MantasV/Procedure_molt5-large")
7
8model.config.max_length = 512
9model.to(torch_device)
10
11#The reactants and products are separated by the bar (|), canonical SMILES format from rdkit
12input = 'Clc1nc(Cl)c2c(n1)CSC2|C1COCCN1>>Clc1nc2c(c(N3CCOCC3)n1)SCC2'
13
14input_enc = tokenizer(input, padding=True, truncation=True, return_tensors='pt').to(torch_device)
15output = model.generate(**input_enc,max_new_tokens=512, num_beams=3, early_stopping=True)
16print(tokenizer.decode(output[0], skip_special_tokens=True))