Spaces:
Runtime error
Runtime error
| import torch | |
| from torch_geometric.data import Dataset | |
| import os | |
| from .context_gen import Reaction_Cluster | |
| import json | |
| from .data_utils import smiles2data, reformat_smiles | |
| from collections import defaultdict | |
| import random | |
| from data_provider.caption_dataset import PretrainCaptionDataset | |
| from data_provider.synthesis_dataset import SynthesisDataset | |
| def format_float_from_string(s): | |
| try: | |
| float_value = float(s) | |
| return f'{float_value:.2f}' | |
| except ValueError: | |
| return s | |
| class MoleculeAbstract(Dataset): | |
| def __init__(self, | |
| root, | |
| rxn_num=1000, | |
| rxn_batch_size=4, | |
| smi_max_len=128, | |
| prompt=None, | |
| disable_graph_cache=False, | |
| disable_graphs=False, | |
| context_style='weighted_rxn', | |
| use_caption_dataset=False, | |
| caption_batch_num=10000, | |
| synthesis_datasetpath=None, | |
| synthesis_batch_num=10000, | |
| reverse_ratio=0.5, | |
| enable_abstract=True, | |
| enable_property=True, | |
| smiles_type='default', | |
| mode='train' | |
| ): | |
| super(MoleculeAbstract, self).__init__(root) | |
| self.root = root | |
| self.rxn_num = rxn_num | |
| self.rxn_batch_size = rxn_batch_size | |
| self.smi_max_len = smi_max_len | |
| self.context_style = context_style | |
| self.tokenizer = None | |
| self.disable_graph_cache = disable_graph_cache | |
| self.disable_graphs = disable_graphs | |
| self.use_caption_dataset = use_caption_dataset | |
| self.smiles_type = smiles_type | |
| if use_caption_dataset: | |
| self.caption_dataset = PretrainCaptionDataset( | |
| os.path.join(root, '../caption_data'), | |
| smi_max_len=smi_max_len, | |
| use_graph=not self.disable_graphs, | |
| disable_graph_cache=disable_graph_cache, | |
| smiles_type=smiles_type, | |
| ) | |
| self.caption_batch_num = caption_batch_num | |
| self.use_synthesis_dataset = bool(synthesis_datasetpath) | |
| if self.use_synthesis_dataset: | |
| self.synthesis_dataset = SynthesisDataset( | |
| synthesis_datasetpath, | |
| 'train', | |
| smi_max_len, | |
| roundrobin_train=True, | |
| use_graph=not disable_graphs, | |
| disable_graph_cache=disable_graph_cache, | |
| smiles_type='default', | |
| ) | |
| self.synthesis_batch_num = synthesis_batch_num | |
| if not self.disable_graphs: | |
| self.mol_graph_map = torch.load(os.path.join(self.root, 'mol_graph_map.pt')) | |
| reaction_filename = 'reactions/reactions_test.json' if (mode=='test') else 'reactions/reactions.json' | |
| if smiles_type=='r_smiles': | |
| reaction_filename = 'reactions/reactions_wRSMILES.json' | |
| self.cluster = Reaction_Cluster(self.root, reaction_filename=reaction_filename, reverse_ratio=reverse_ratio) | |
| self.reload_data_list() | |
| self.abstract_max_len = 10240 | |
| self.property_max_len = 10240 | |
| self.enable_abstract = enable_abstract | |
| self.enable_property = enable_property | |
| def get(self, index): | |
| return self.__getitem__(index) | |
| def len(self): | |
| return len(self) | |
| def __len__(self): | |
| data_len = len(self.data_list) | |
| if self.use_caption_dataset: | |
| data_len += len(self.caption_index_list) | |
| if self.use_synthesis_dataset: | |
| data_len += len(self.synthesis_index_list) | |
| return data_len | |
| def reload_data_list(self): | |
| k = self.rxn_batch_size | |
| if self.context_style == 'weighted_rxn': | |
| self.data_list = self.cluster(self.rxn_num, k=k) | |
| elif self.context_style == 'uniform_rxn': | |
| self.data_list = self.cluster.generate_batch_uniform_rxn(self.rxn_num, k=k) | |
| elif self.context_style == 'uniform_mol': | |
| self.data_list = self.cluster.generate_batch_uniform_mol(self.rxn_num, k=k) | |
| elif self.context_style == 'single_mol': | |
| self.data_list = self.cluster.generate_batch_single(self.rxn_num) | |
| elif self.context_style == 'hybrid': | |
| self.data_list = self.cluster(self.rxn_num//2, k=k) | |
| self.data_list += self.cluster.generate_batch_uniform_mol(self.rxn_num//2, k=k) | |
| else: | |
| raise NotImplementedError | |
| if self.use_caption_dataset: | |
| assert self.caption_batch_num*k <= len(self.caption_dataset) | |
| caption_index_list = random.sample(range(len(self.caption_dataset)), self.caption_batch_num*k) | |
| self.caption_index_list = [caption_index_list[i*k:(i+1)*k] for i in range(self.caption_batch_num)] | |
| else: | |
| self.caption_index_list = [] | |
| if self.use_synthesis_dataset: | |
| if self.synthesis_dataset.roundrobin_train: | |
| self.synthesis_dataset.reload_data() | |
| assert self.synthesis_batch_num <= len(self.synthesis_dataset) | |
| self.synthesis_index_list = random.sample(range(len(self.synthesis_dataset)), self.synthesis_batch_num) | |
| else: | |
| self.synthesis_index_list = [] | |
| def make_prompt(self, mol_batch, smi_max_len=128): | |
| mol_prompt_list, text_prompt_list = [], [] | |
| last_role = None | |
| for mol_dict in mol_batch: | |
| smiles = mol_dict['canon_smiles'] | |
| if self.smiles_type=='r_smiles': | |
| if 'r_smiles' in mol_dict: | |
| smiles = mol_dict['r_smiles'] | |
| # else: | |
| # smiles = reformat_smiles(smiles, smiles_type='restricted') | |
| else: | |
| smiles = reformat_smiles(smiles, smiles_type=self.smiles_type) | |
| mol_prompt = f'[START_SMILES]{smiles[:smi_max_len]}[END_SMILES]. ' | |
| if 'role' in mol_dict: | |
| role = { | |
| 'REACTANT': 'Reactant', | |
| 'CATALYST': 'Catalyst', | |
| 'SOLVENT': 'Solvent', | |
| 'PRODUCT': 'Product', | |
| }[mol_dict['role']] | |
| if last_role != role: | |
| mol_prompt = f'{role}: {mol_prompt}' | |
| last_role = role | |
| text_prompt = self.make_abstract(mol_dict) | |
| mol_prompt_list.append(mol_prompt) | |
| text_prompt_list.append(text_prompt) | |
| return mol_prompt_list, text_prompt_list | |
| def make_abstract(self, mol_dict): | |
| prompt = '' | |
| if self.enable_abstract and 'abstract' in mol_dict: | |
| abstract_string = mol_dict['abstract'][:self.abstract_max_len] | |
| prompt += f'[Abstract] {abstract_string} ' | |
| if self.enable_property: | |
| property_string = '' | |
| property_dict = mol_dict['property'] if 'property' in mol_dict else {} | |
| for property_key in ['Experimental Properties', 'Computed Properties']: | |
| if not property_key in property_dict: | |
| continue | |
| for key, value in property_dict[property_key].items(): | |
| if isinstance(value, float): | |
| key_value_string = f'{key}: {value:.2f}; ' | |
| elif isinstance(value, str): | |
| float_value = format_float_from_string(value) | |
| key_value_string = f'{key}: {float_value}; ' | |
| else: | |
| key_value_string = f'{key}: {value}; ' | |
| if len(property_string+key_value_string) > self.property_max_len: | |
| break | |
| property_string += key_value_string | |
| if property_string: | |
| property_string = property_string[:self.property_max_len] | |
| prompt += f'[Properties] {property_string}. ' | |
| return prompt | |
| def get_caption_data(self, index): | |
| caption_index = self.caption_index_list[index] | |
| graph_list, mol_prompt_list, text_prompt_list = [], [], [] | |
| for idx in caption_index: | |
| graph_item, text, smiles_prompt = self.caption_dataset[idx] | |
| graph_list.append(graph_item) | |
| mol_prompt_list.append(smiles_prompt) | |
| text_prompt_list.append(text) | |
| return graph_list, mol_prompt_list, text_prompt_list | |
| def get_synthesis_data(self, index): | |
| synthesis_index = self.synthesis_index_list[index] | |
| _, graph_list, output_text, input_text = self.synthesis_dataset[synthesis_index] | |
| return graph_list, [input_text], [output_text] | |
| def __getitem__(self, index): | |
| if index < len(self.data_list): | |
| mol_batch = self.data_list[index] | |
| elif index < len(self.data_list)+len(self.caption_index_list): | |
| assert self.use_caption_dataset | |
| return self.get_caption_data(index-len(self.data_list)) | |
| else: | |
| assert self.use_synthesis_dataset | |
| return self.get_synthesis_data(index-(len(self.data_list)+len(self.caption_index_list))) | |
| graph_list = [] | |
| for mol_dict in mol_batch: | |
| smiles = mol_dict['canon_smiles'] | |
| if self.disable_graphs: | |
| graph_item = None | |
| else: | |
| if self.disable_graph_cache: | |
| graph_item = smiles2data(smiles) | |
| else: | |
| assert smiles in self.mol_graph_map | |
| graph_item = self.mol_graph_map[smiles] | |
| graph_list.append(graph_item) | |
| mol_prompt_list, text_prompt_list = self.make_prompt(mol_batch, smi_max_len=self.smi_max_len) | |
| return graph_list, mol_prompt_list, text_prompt_list | |