Views
No views yet
1from transformers import T5Tokenizer, T5ForConditionalGeneration
2
3tokenizer = T5Tokenizer.from_pretrained("QizhiPei/biot5-base-mol2text", model_max_length=512)
4model = T5ForConditionalGeneration.from_pretrained('QizhiPei/biot5-base-mol2text')
5
6task_definition = 'Definition: You are given a molecule SELFIES. Your job is to generate the molecule description in English that fits the molecule SELFIES.\n\n'
7selfies_input = '[C][C][Branch1][C][O][C][C][=Branch1][C][=O][C][=Branch1][C][=O][O-1]'
8task_input = f'Now complete the following example -\nInput: <bom>{selfies_input}<eom>\nOutput: '
9
10model_input = task_definition + task_input
11input_ids = tokenizer(model_input, return_tensors="pt").input_ids
12
13generation_config = model.generation_config
14generation_config.max_length = 512
15generation_config.num_beams = 1
16
17outputs = model.generate(input_ids, generation_config=generation_config)
18
19print(tokenizer.decode(outputs[0], skip_special_tokens=True))