# Copyright 2024 the Llamole Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel, PreTrainedTokenizerBase from transformers.utils import ModelOutput from transformers.generation.utils import LogitsProcessorList, GenerationConfig from huggingface_hub import snapshot_download from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union import os import json import time from dataclasses import dataclass from typing import Union, Tuple, Optional from .loader import load_language_model, load_tokenizer from .loader import load_graph_decoder, load_graph_predictor, load_graph_encoder from ..extras.constants import NO_LABEL_INDEX, IGNORE_INDEX, BOND_INDEX from .planner import molstar from rdkit import Chem from torch_geometric.data import Data from torch_geometric.data import Batch as PyGBatch from torch_geometric.utils import remove_isolated_nodes # Save configuration def convert_to_dict(obj): if isinstance(obj, (int, float, str, bool, type(None))): return obj elif isinstance(obj, (list, tuple)): return [convert_to_dict(item) for item in obj] elif isinstance(obj, dict): return {k: convert_to_dict(v) for k, v in obj.items()} elif hasattr(obj, "__dict__"): return { k: convert_to_dict(v) for k, v in obj.__dict__.items() if not k.startswith("_") } else: return str(obj) # Convert any other objects to string @dataclass class GraphLMOutput(ModelOutput): loss: Optional[torch.FloatTensor] = None logits: torch.FloatTensor = None last_hidden_state: Optional[torch.FloatTensor] = None past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None attentions: Optional[Tuple[torch.FloatTensor]] = None additional_log_info: Optional[Dict[str, float]] = None class GraphLLMForCausalMLM(PreTrainedModel): def __init__( self, model_args, finetuning_args, data_args, language_model, graph_decoder, graph_predictor, graph_encoder, token_id_dict, tokenizer, ): super().__init__(language_model.config) self.language_model = language_model self.graph_decoder = graph_decoder self.graph_predictor = graph_predictor self.graph_encoder = graph_encoder self.token_id_dict = token_id_dict self.num_body_tokens = data_args.learned_query_size self.loss_weight_lm = finetuning_args.loss_weight_lm self.loss_weight_design = finetuning_args.loss_weight_design self.loss_weight_retro = finetuning_args.loss_weight_retro self.model_args = model_args self.finetuning_args = finetuning_args self.data_args = data_args self.tokenizer = tokenizer # Initialize weights and apply final processing self.post_init() @classmethod def from_pretrained( cls, tokenizer: PreTrainedTokenizerBase, model_args, data_args, training_args, finetuning_args, load_adapter=False, add_valuehead=False, ): if load_adapter: if model_args.adapter_name_or_path is None: raise ValueError("Please specify the adapter_name_or_path when load_adapter is True.") if len(model_args.adapter_name_or_path) != 1: raise ValueError("Only one adapter is supported at a time.") adapter_path = model_args.adapter_name_or_path[0] if not os.path.exists(os.path.join(adapter_path, "adapter_config.json")): # Download from HuggingFace adapter_name = os.path.basename(adapter_path) valid_adapters = [ "Llama-3.1-8B-Instruct-Adapter", "Qwen2-7B-Instruct-Adapter", "Mistral-7B-Instruct-v0.3-Adapter" ] if adapter_name not in valid_adapters: raise ValueError(f"Invalid adapter name. Supported adapters are: {', '.join(valid_adapters)}") repo_id = f"liuganghuggingface/Llamole-{adapter_name}" print(f"Downloading adapter {adapter_name} from HuggingFace repo: {repo_id}") try: # Download all files including subfolders to the adapter_path snapshot_download( repo_id=repo_id, local_dir=adapter_path, local_dir_use_symlinks=False, ignore_patterns=["*.md", "*.txt"] # Optionally ignore certain file types ) print(f"Successfully downloaded all adapter files to {adapter_path}") except Exception as e: raise RuntimeError(f"Failed to download adapter files: {str(e)}") language_model = load_language_model( tokenizer, model_args, finetuning_args, training_args.do_train, add_valuehead, ) device = next(language_model.parameters()).device graph_decoder = load_graph_decoder( model_args, path=model_args.graph_decoder_path, device=device, ) graph_predictor = load_graph_predictor( model_args, path=model_args.graph_predictor_path, device=device, ) graph_encoder = load_graph_encoder( model_args, path=model_args.graph_encoder_path, device=device, ) if ( getattr(language_model, "is_quantized", False) and not training_args.do_train ): setattr( language_model, "_hf_peft_config_loaded", True ) # hack here: make model compatible with prediction token_id_dict = {} for elem in model_args.new_special_tokens: if isinstance(elem, str) and len(elem) != 0: elem_token_ids = tokenizer.encode(elem, add_special_tokens=False) token_id_dict[elem] = elem_token_ids[0] model = cls( model_args=model_args, finetuning_args=finetuning_args, data_args=data_args, language_model=language_model, graph_decoder=graph_decoder, graph_predictor=graph_predictor, graph_encoder=graph_encoder, token_id_dict=token_id_dict, tokenizer=tokenizer, ) graph_to_lm_connector = nn.Sequential( nn.Linear(graph_encoder.hidden_size, language_model.config.hidden_size), nn.SiLU(), ) # Language Model to Graph Decoder connector lm_to_graph_decoder = nn.Sequential( nn.Linear(language_model.config.hidden_size, graph_decoder.text_input_size), nn.SiLU(), ) # Language Model to Graph Predictor connector lm_to_graph_predictor = nn.Sequential( nn.Linear( language_model.config.hidden_size, graph_predictor.text_input_size ), nn.SiLU(), ) for param in graph_to_lm_connector.parameters(): if ( param.dtype == torch.float32 and model_args.compute_dtype != torch.float32 ): param.data = param.data.to(model_args.compute_dtype) for param in lm_to_graph_decoder.parameters(): if ( param.dtype == torch.float32 and model_args.compute_dtype != torch.float32 ): param.data = param.data.to(model_args.compute_dtype) for param in lm_to_graph_predictor.parameters(): if ( param.dtype == torch.float32 and model_args.compute_dtype != torch.float32 ): param.data = param.data.to(model_args.compute_dtype) # Check if connector path is provided and load if available if load_adapter: if ( hasattr(model_args, "graph_lm_connector_path") and model_args.graph_lm_connector_path ): connector_path = model_args.graph_lm_connector_path graph_to_lm_connector.load_state_dict( torch.load( os.path.join(connector_path, "graph_to_lm_connector.pt"), map_location=device, weights_only=True, ) ) lm_to_graph_decoder.load_state_dict( torch.load( os.path.join(connector_path, "lm_to_graph_decoder.pt"), map_location=device, weights_only=True, ) ) lm_to_graph_predictor.load_state_dict( torch.load( os.path.join(connector_path, "lm_to_graph_predictor.pt"), map_location=device, weights_only=True, ) ) else: raise ValueError(f"Connector should be automatically downloaded with the adapter. Please manually download to the path {connector_path}") model.graph_to_lm_connector = graph_to_lm_connector model.lm_to_graph_decoder = lm_to_graph_decoder model.lm_to_graph_predictor = lm_to_graph_predictor model.graph_to_lm_connector.to(device) model.lm_to_graph_decoder.to(device) model.lm_to_graph_predictor.to(device) return model def to(self, device): super().to(device) self.language_model.to(device) self.graph_decoder.to(device) self.graph_predictor.to(device) self.graph_encoder.to(device) self.graph_to_lm_connector.to(device) self.lm_to_graph_decoder.to(device) self.lm_to_graph_predictor.to(device) return self def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, molecule_graphs: Optional[PyGBatch] = None, molecule_properties: Optional[torch.FloatTensor] = None, design_graphs: Optional[PyGBatch] = None, retro_labels: Optional[torch.LongTensor] = None, retro_product_graphs: Optional[PyGBatch] = None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = True, return_dict: Optional[bool] = None, ) -> Union[Tuple, GraphLMOutput]: return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) mol_token_id = self.token_id_dict[""] design_start_token_id = self.token_id_dict[""] retro_start_token_id = self.token_id_dict[""] # PeftModelForCausalLM -> LlamaForCausalLM -> LlamaModel base_llm = self.language_model.model.model inputs_embeds = base_llm.embed_tokens(input_ids) mol_positions = (input_ids == mol_token_id).nonzero() mol_embeds = self.graph_encoder( molecule_graphs.x, molecule_graphs.edge_index, molecule_graphs.edge_attr, molecule_graphs.batch, ) mol_embeds = self.graph_to_lm_connector(mol_embeds) assert ( mol_positions.shape[0] == mol_embeds.shape[0] ), f"Number of molecule tokens ({mol_positions.shape[0]}) does not match number of molecule embeddings ({mol_embeds.shape[0]})" inputs_embeds[mol_positions[:, 0], mol_positions[:, 1]] = mol_embeds.to( inputs_embeds.dtype ) lm_outputs = self.language_model( input_ids=None, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, labels=labels, ) lm_loss = lm_outputs.loss lm_hidden_states = lm_outputs.hidden_states[-1] design_loss = 0 if design_graphs is not None: design_start_positions = (input_ids == design_start_token_id).nonzero() design_body_start = design_start_positions[:, 1] + 1 design_body_indices = design_body_start.unsqueeze(1) + torch.arange( self.num_body_tokens, device=input_ids.device ) design_hidden = lm_hidden_states[ design_start_positions[:, 0].unsqueeze(1), design_body_indices[:, 1] ].mean(dim=1) if design_start_positions.numel() > 0: design_hidden = self.lm_to_graph_decoder(design_hidden) design_loss = self.graph_decoder( design_graphs.x, design_graphs.edge_index, design_graphs.edge_attr, design_graphs.batch, molecule_properties, design_hidden, NO_LABEL_INDEX, ) # Process retro labels retro_loss = 0 if retro_labels is not None: # Get retro start positions for valid retro labels: (batch, step) retro_start_positions = (input_ids == retro_start_token_id).nonzero() retro_labels = retro_labels[retro_labels != IGNORE_INDEX] valid_retro_mask = retro_labels != NO_LABEL_INDEX retro_start_positions = retro_start_positions[valid_retro_mask] retro_labels = retro_labels[valid_retro_mask] if len(retro_labels) > 0: # Get the query hidden states for each retro prediction retro_body_start = retro_start_positions[:, 1] + 1 retro_body_indices = retro_body_start.unsqueeze(1) + torch.arange( self.num_body_tokens, device=input_ids.device ) retro_hidden = lm_hidden_states[ retro_start_positions[:, 0].unsqueeze(1), retro_body_indices ].mean(dim=1) # Prepare graph inputs retro_product_graphs = retro_product_graphs[ valid_retro_mask.nonzero().view(-1) ] retro_product_graphs = PyGBatch.from_data_list(retro_product_graphs) # Transform hidden states and make predictions retro_hidden = self.lm_to_graph_predictor(retro_hidden) retro_pred = self.graph_predictor( retro_product_graphs.x, retro_product_graphs.edge_index, retro_product_graphs.edge_attr, retro_product_graphs.batch, retro_hidden, ) retro_loss = F.cross_entropy( retro_pred, retro_labels, ) total_loss = ( self.loss_weight_lm * lm_loss + self.loss_weight_design * retro_loss + self.loss_weight_retro * retro_loss ) if not return_dict: output = (lm_outputs.logits,) + lm_outputs[1:] return ((total_loss,) + output) if total_loss is not None else output return GraphLMOutput( loss=total_loss, logits=lm_outputs.logits, past_key_values=lm_outputs.past_key_values, hidden_states=lm_outputs.hidden_states, attentions=lm_outputs.attentions, ) def save_pretrained( self, save_directory: Union[str, os.PathLike], is_main_process: bool = True, state_dict: Optional[dict] = None, save_function: Callable = torch.save, push_to_hub: bool = False, max_shard_size: Union[int, str] = "5GB", safe_serialization: bool = True, variant: Optional[str] = None, token: Optional[Union[str, bool]] = None, save_peft_format: bool = True, save_graph_modules: bool = False, **kwargs, ): """ Save the model and its configuration file to a directory. """ if os.path.isfile(save_directory): raise ValueError( f"Provided path ({save_directory}) should be a directory, not a file" ) os.makedirs(save_directory, exist_ok=True) # Save language model language_model_path = os.path.join(save_directory) self.language_model.save_pretrained( language_model_path, is_main_process=is_main_process, state_dict=state_dict, save_function=save_function, push_to_hub=False, # set to false max_shard_size=max_shard_size, safe_serialization=safe_serialization, variant=variant, token=token, save_peft_format=save_peft_format, ) if save_graph_modules: # Save graph models graph_models = { "graph_decoder": self.graph_decoder, "graph_predictor": self.graph_predictor, "graph_encoder": self.graph_encoder, } for name, model in graph_models.items(): model_path = os.path.join(save_directory, name) model.save_pretrained(model_path) # Save additional components additional_components = { "graph_to_lm_connector": self.graph_to_lm_connector, "lm_to_graph_decoder": self.lm_to_graph_decoder, "lm_to_graph_predictor": self.lm_to_graph_predictor, } connector_path = os.path.join(save_directory, "connector") for name, component in additional_components.items(): os.makedirs(connector_path, exist_ok=True) component_path = os.path.join(connector_path, f"{name}.pt") torch.save(component.state_dict(), component_path) config_dict = { "model_args": convert_to_dict(self.model_args), "finetuning_args": convert_to_dict(self.finetuning_args), "data_args": convert_to_dict(self.data_args), "token_id_dict": self.token_id_dict, "num_body_tokens": self.num_body_tokens, "loss_weight_lm": self.loss_weight_lm, "loss_weight_design": self.loss_weight_design, "loss_weight_retro": self.loss_weight_retro, } config_path = os.path.join(save_directory, "graphllm_config.json") with open(config_path, "w") as f: json.dump(config_dict, f, indent=2) # Push to hub if required if push_to_hub: raise NotImplementedError("Push to hub not implemented yet") def add_special_body_tokens( self, input_ids: torch.LongTensor, body_token_id: int, num_body_tokens: int, start_token_id: Optional[int] = None, ) -> torch.LongTensor: batch_size, seq_length = input_ids.shape start_len = 1 if start_token_id is not None else 0 if seq_length < num_body_tokens + start_len: seq_length = seq_length + num_body_tokens + start_len # Create a tensor to hold start positions for each batch item start_positions = torch.full( (batch_size,), seq_length - start_len - num_body_tokens, device=input_ids.device, ) # Calculate how many tokens to keep from the original input tokens_to_keep = seq_length - num_body_tokens # Find start positions if start_token_id is not None: start_pos_rows, start_pos_cols = (input_ids == start_token_id).nonzero( as_tuple=True ) for row, col in zip(start_pos_rows, start_pos_cols): start_positions[row] = col tokens_to_keep = seq_length - num_body_tokens - 1 # Create body tokens body_tokens = torch.full( (batch_size, num_body_tokens), body_token_id, device=input_ids.device ) # Create new input_ids with left padding new_input_ids = torch.full( (batch_size, seq_length), self.tokenizer.eos_token_id, device=input_ids.device, ) for i in range(batch_size): start_pos = start_positions[i] # Keep the rightmost tokens_to_keep tokens before the start token keep_start = max(0, start_pos - tokens_to_keep) if start_token_id is not None: new_input_ids[ i, -(num_body_tokens + 1 + (start_pos - keep_start)) : ] = torch.cat( [ input_ids[i, keep_start:start_pos], torch.LongTensor([start_token_id]).to(input_ids.device), body_tokens[i], ] ) else: new_input_ids[ i, -(num_body_tokens + 1 + (start_pos - keep_start)) : ] = torch.cat([input_ids[i, keep_start:start_pos], body_tokens[i]]) return new_input_ids @torch.no_grad() def design_molecule( self, input_ids: torch.LongTensor, attention_mask: torch.FloatTensor, molecule_properties: Optional[torch.FloatTensor] = None, molecule_graphs: Optional[PyGBatch] = None, rollback: bool = False, **kwargs, ) -> List[Optional[str]]: design_start_token_id = self.token_id_dict[""] design_body_token_id = self.token_id_dict[""] # 1. Generate molecular design analysis if molecule_graphs is None: analysis_tokens = self.language_model.generate( inputs=input_ids, attention_mask=attention_mask, **kwargs, ) analysis_tokens = analysis_tokens[:, input_ids.shape[1] :] else: mol_token_id = self.token_id_dict[""] base_llm = self.language_model.model inputs_embeds = base_llm.embed_tokens(input_ids) mol_positions = (input_ids == mol_token_id).nonzero() mol_embeds = self.graph_encoder( molecule_graphs.x, molecule_graphs.edge_index, molecule_graphs.edge_attr, molecule_graphs.batch, ) mol_embeds = self.graph_to_lm_connector(mol_embeds) assert ( mol_positions.shape[0] == mol_embeds.shape[0] ), f"Number of molecule tokens ({mol_positions.shape[0]}) does not match number of molecule embeddings ({mol_embeds.shape[0]})" inputs_embeds[mol_positions[:, 0], mol_positions[:, 1]] = mol_embeds.to( inputs_embeds.dtype ) analysis_tokens = self.language_model.generate( attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs, ) # no input # 2. Add special tokens for design body design_input_ids = self.add_special_body_tokens( analysis_tokens, design_body_token_id, self.num_body_tokens, start_token_id=design_start_token_id, ) design_input_ids = torch.cat([input_ids, design_input_ids], dim=1) # 3. Get LLM embeddings for design body lm_outputs = self.language_model( input_ids=design_input_ids, attention_mask=torch.ones_like(design_input_ids), output_hidden_states=True, return_dict=True, ) lm_hidden_states = lm_outputs.hidden_states[-1] design_hidden = lm_hidden_states[:, -self.num_body_tokens :].mean(dim=1) # 4. Generate molecules using graph decoder design_hidden = self.lm_to_graph_decoder(design_hidden) molecule_properties = molecule_properties.type_as(design_hidden) smiles_list = self.graph_decoder.generate( molecule_properties, design_hidden, NO_LABEL_INDEX, ) # Handle None values in smiles_list if rollback and None in smiles_list: smiles_list = self.design_rollback(design_input_ids, smiles_list, **kwargs) return analysis_tokens, smiles_list def design_rollback( self, analysis_tokens: torch.LongTensor, smiles_list: List[Optional[str]], **kwargs, ) -> List[Optional[str]]: rollback_token_id = self.token_id_dict.get("") rollback_end_token_id = self.token_id_dict.get("") none_indices = [i for i, smiles in enumerate(smiles_list) if smiles is None] if not none_indices: return smiles_list # No None values, return original list # Get corresponding analysis tokens for None indices none_indices = torch.LongTensor(none_indices) rollback_analysis_tokens = analysis_tokens[none_indices] # Add rollback token to the end of each analysis token sequence rollback_input_ids = self.add_special_body_tokens( rollback_analysis_tokens, rollback_token_id, 1, ) if "max_new_tokens" in kwargs: kwargs["max_new_tokens"] *= 2 # Generate new tokens new_tokens = self.language_model.generate( inputs=rollback_input_ids, attention_mask=torch.ones_like(rollback_input_ids), **kwargs, ) # Process and decode new tokens new_smiles = [] for seq in new_tokens[:, rollback_input_ids.shape[1] :]: decoded_seq = self.tokenizer.decode(seq, skip_special_tokens=False) end_smiles_pos = decoded_seq.find( self.tokenizer.decode([rollback_end_token_id]) ) if end_smiles_pos != -1: # If end token is found, truncate the sequence new_smiles.append(decoded_seq[:end_smiles_pos].strip()) else: # If end token is not found, append None new_smiles.append(None) # Update smiles_list with new decoded tokens for i, new_smiles_str in zip(none_indices, new_smiles): smiles_list[i] = new_smiles_str return smiles_list def smiles_to_graph(self, smiles: str) -> Optional[Data]: mol = Chem.MolFromSmiles(smiles) if mol is None: print(f"Invalid SMILES string: {smiles}") return None type_idx = [] for atom in mol.GetAtoms(): if atom.GetAtomicNum() != 1: # Exclude hydrogen atoms type_idx.append( 119 - 2 if atom.GetSymbol() == "*" else atom.GetAtomicNum() - 2 ) x = torch.LongTensor(type_idx) num_nodes = x.size(0) # Initialize edge_index and edge_attr as empty tensors edge_index = torch.empty((2, 0), dtype=torch.long) edge_attr = torch.empty((0,), dtype=torch.long) # Only process bonds if they exist if mol.GetNumBonds() > 0: bond_src = [] bond_dst = [] bond_type = [] for bond in mol.GetBonds(): start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx() # Exclude bonds involving hydrogen atoms if mol.GetAtomWithIdx(start).GetAtomicNum() != 1 and mol.GetAtomWithIdx(end).GetAtomicNum() != 1: bond_src.extend([start, end]) bond_dst.extend([end, start]) bond_type.extend([BOND_INDEX.get(bond.GetBondType(), 1)] * 2) if bond_src: # Only create edge_index and edge_attr if there are valid bonds edge_index = torch.tensor([bond_src, bond_dst], dtype=torch.long) edge_attr = torch.tensor(bond_type, dtype=torch.long) # Create PyG Data object data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, num_nodes=num_nodes) return data def retrosynthesize_rollback(self, input_ids, design_text, smiles, **kwargs): input_text = f"{design_text} To synthesize {smiles}, follow these procedures: " input_tokens = self.tokenizer.encode( input_text, add_special_tokens=False, return_tensors="pt" ) input_tokens = input_tokens.to(self.device) if "max_new_tokens" in kwargs: kwargs["max_new_tokens"] = 256 # Generate tokens generated_tokens = self.language_model.generate( inputs=input_tokens, **kwargs, ) generated_tokens = generated_tokens[:, input_tokens.shape[1] :] generated_tokens = generated_tokens.cpu().squeeze().tolist() new_input_text = f"To synthesize {smiles}, follow these procedures: " new_input_tokens = self.tokenizer.encode(new_input_text) generated_tokens = new_input_tokens + generated_tokens return generated_tokens def one_step_reaction( self, product_smiles, input_ids, design_text, molecule_graphs, topk, **kwargs, ): # 1. Generate retrosynthesis analysis retro_start_token_id = self.token_id_dict[""] retro_body_token_id = self.token_id_dict[""] mol_token_id = self.token_id_dict[""] input_text = f"{design_text} To synthesize , follow these procedures: " prompt_tokens = self.tokenizer.encode( input_text, add_special_tokens=False, return_tensors="pt" ) prompt_tokens = prompt_tokens.to(self.device) # Combine input_ids with new_prompt_tokens if input_ids is provided if input_ids is not None and molecule_graphs is not None: input_ids = input_ids.view(1, -1) prompt_tokens = torch.cat([input_ids, prompt_tokens], dim=-1) base_llm = self.language_model.model inputs_embeds = base_llm.embed_tokens(prompt_tokens) product_graph = self.smiles_to_graph(product_smiles) if product_graph is None: return { "reactants": [], "scores": [], "templates": [], "analysis": self.tokenizer.encode( "Invalid product SMILES", add_special_tokens=False ), } product_graph.to(self.device) if input_ids is not None and molecule_graphs is not None: all_graphs = PyGBatch.from_data_list(molecule_graphs.to_data_list() + [product_graph]) else: all_graphs = PyGBatch.from_data_list([product_graph]) mol_embeds = self.graph_encoder( all_graphs.x, all_graphs.edge_index, all_graphs.edge_attr, all_graphs.batch, ) mol_embeds = self.graph_to_lm_connector(mol_embeds) mol_positions = (prompt_tokens == mol_token_id).nonzero() assert ( mol_positions.shape[0] == mol_embeds.shape[0] ), f"Number of molecule tokens ({mol_positions.shape[0]}) does not match number of molecule embeddings ({mol_embeds.shape[0]})" inputs_embeds[mol_positions[:, 0], mol_positions[:, 1]] = mol_embeds.to( inputs_embeds.dtype ) attention_mask = torch.ones_like(prompt_tokens) if "max_new_tokens" in kwargs: kwargs["max_new_tokens"] = 512 analysis_tokens = self.language_model.generate( attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs, ) # 2. Encode analysis with query tokens retro_input_ids = self.add_special_body_tokens( analysis_tokens, retro_body_token_id, self.num_body_tokens, start_token_id=retro_start_token_id, ) # Get LLM embeddings for retro body lm_outputs = self.language_model( input_ids=retro_input_ids, attention_mask=torch.ones_like(retro_input_ids), output_hidden_states=True, return_dict=True, ) lm_hidden_states = lm_outputs.hidden_states[-1] retro_hidden = lm_hidden_states[:, -self.num_body_tokens :].mean(dim=1) retro_hidden = self.lm_to_graph_predictor(retro_hidden) # 3. Sample retrosynthetic templates reactants, scores, templates = self.graph_predictor.sample_templates( product_graph, retro_hidden, product_smiles, topk ) # 4. Adjust the input part from the generated tokens analysis_tokens = analysis_tokens.cpu().squeeze().tolist() input_text = f"To synthesize {product_smiles}, follow these procedures: " new_input_tokens = self.tokenizer.encode(input_text) analysis_tokens = new_input_tokens + analysis_tokens return { "reactants": reactants, "scores": scores, "templates": templates, "analysis": analysis_tokens, } @torch.no_grad() def estimate_synthesis_complexity( self, smiles: str, input_ids=None, reaction=None, molecule_cost_weight: float = 0, language_cost_weight: float = 1, reference_tokens: Optional[torch.LongTensor] = None, ): cost = 0 if molecule_cost_weight is not None and molecule_cost_weight > 0: mol_cost = self.graph_predictor.estimate_cost(smiles) cost += mol_cost * molecule_cost_weight if language_cost_weight is not None and language_cost_weight > 0: language_cost = 0 if reaction is None: message_content = f""" Estimate remaining steps for the target {smiles} consider the following factors:: 1. Intermediate complexity 2. Reagent availability 3. Side reactions 4. Stereochemistry challenges""" else: step = reaction.depth + 1 template = reaction.template # analysis_tokens = reaction.analysis_tokens reactants = reaction.children reactants = ", ".join([r.mol for r in reactants]) message_content = f""" Estimate remaining steps for the target {smiles} given the following parameters: Current step {step}, Current template: {template}, Reactants: {reactants}. Consider the following factors: 1. Intermediate complexity 2. Reagent availability 3. Side reactions 4. Stereochemistry challenges""" # Create the messages list for the chat template messages = [{"role": "user", "content": message_content}] # Apply the chat template chat_text = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) answers = [ "All readily available", "Some commercial, some need 1-2 steps", "Mix of commercial and multi-step synthesis", "Mostly require complex synthesis", "All require extensive multi-step synthesis", ] answer_costs = [0, 1, 2.5, 4.5, 7] answer_messages = [ [ { "role": "user", "content": "Estimate the synthesis complexity:", }, {"role": "assistant", "content": answer}, ] for answer in answers ] answer_chat_texts = [ self.tokenizer.apply_chat_template( msg, tokenize=False, add_generation_prompt=False ) for msg in answer_messages ] # Encode chat texts input_ids = self.tokenizer.encode(chat_text, return_tensors="pt").to( self.device ) answer_tokens = [ self.tokenizer.encode(text) for text in answer_chat_texts ] # Get logits from the language model outputs = self.language_model(input_ids) logits = outputs.logits[:, -1, :] # Calculate softmax probabilities for each answer answer_logits = torch.stack( [logits[:, tokens].mean(dim=1) for tokens in answer_tokens] ) probs = torch.nn.functional.softmax(answer_logits, dim=0) language_cost = ( (probs * torch.tensor(answer_costs, device=probs.device)) .sum() .item() ) language_cost = language_cost * language_cost_weight cost += language_cost return cost @torch.no_grad() def retrosynthesize( self, input_ids: torch.LongTensor, smiles: Optional[str] = None, molecule_graphs: Optional[PyGBatch] = None, expansion_topk: int = 50, iterations: int = 100, starting_mols: Optional[List[str]] = None, molecule_cost_weight: float = 0, language_cost_weight: float = 1, max_planning_time: int = 300, rollback: bool = True, design_text: Optional[str] = None, **kwargs, ) -> Dict[str, Any]: # Initialize variables target_smiles = None success = False reaction_list = None template_list = None analysis_tokens_list = None route_length = None total_time = 0.0 cost = None # Handle starting molecules if starting_mols is None: if self.graph_predictor.available is None: raise ValueError( "No starting molecules provided and no available starting molecules found." ) starting_mols = self.graph_predictor.available["smiles"].tolist() # Handle case when no SMILES is provided if smiles is None and rollback: generated_tokens = self.retrosynthesize_rollback(input_ids, design_text, None, **kwargs) return self._create_failure_result(None, generated_tokens) # Preprocess SMILES target_smiles = smiles.replace("*", "[H]") if "*" in smiles else smiles # Check validity and handle rollback if necessary if not self.graph_decoder.check_valid(target_smiles) and rollback: generated_tokens = self.retrosynthesize_rollback( input_ids, design_text, target_smiles, **kwargs ) return self._create_failure_result(target_smiles, generated_tokens) # Perform retrosynthesis t0 = time.time() def expand_fn(s): return self.one_step_reaction( s, input_ids=input_ids, design_text=design_text, molecule_graphs=molecule_graphs, topk=expansion_topk, **kwargs ) def value_fn(s, r): return self.estimate_synthesis_complexity( s, input_ids, r, molecule_cost_weight, language_cost_weight ) if target_smiles is None: return self._create_failure_result(None) success, best_route, iterations = molstar( target_mol=target_smiles, target_mol_id=0, starting_mols=starting_mols, expand_fn=expand_fn, value_fn=value_fn, iterations=iterations, max_time=max_planning_time, ) total_time = time.time() - t0 # Handle successful retrosynthesis if success: reaction_list, template_list, cost, analysis_tokens_list = best_route.get_reaction_list() route_length = best_route.length # Handle failed retrosynthesis with rollback elif rollback: generated_tokens = self.retrosynthesize_rollback( input_ids, design_text, target_smiles, **kwargs ) return self._create_failure_result(target_smiles, generated_tokens) # Prepare and return result return { "target": target_smiles, "success": success, "time": total_time, "reaction_list": reaction_list, "cost": cost, "templates": template_list, "analysis_tokens": analysis_tokens_list, "route_length": route_length, } def _create_failure_result( self, target_smiles: Optional[str], generated_tokens: Optional[Union[torch.Tensor, list]] = None, ) -> Dict[str, Any]: return { "target": target_smiles, "success": False, "time": 0.0, "reaction_list": None, "cost": None, "templates": None, "analysis_tokens": ( generated_tokens if generated_tokens is not None else "" ), "route_length": None, } @torch.no_grad() def generate( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, molecule_properties: Optional[torch.FloatTensor] = None, molecule_graphs: Optional[PyGBatch] = None, rollback: bool = False, starting_mols: Optional[List[str]] = None, expansion_topk: int = 50, iterations: int = 100, molecule_cost_weight: float = 0, language_cost_weight: float = 1, do_molecular_design: Optional[bool] = True, do_retrosynthesis: bool = True, input_smiles_list: Optional[List[str]] = None, max_planning_time: int = 30, design_text_list: Optional[List[str]] = None, **kwargs, ) -> Dict: if attention_mask is None: attention_mask = input_ids.new_ones(input_ids.shape) all_info_dict = { "token_lists": [], "text_lists": [], "design_analysis_tokens": None, "smiles_list": None, "retro_plan_dict": None, } # Molecular design if do_molecular_design is True: design_analysis_tokens, smiles_list = self.design_molecule( input_ids, attention_mask, molecule_properties, molecule_graphs, rollback, **kwargs, ) all_info_dict["design_analysis_tokens"] = design_analysis_tokens.cpu() all_info_dict["smiles_list"] = smiles_list elif input_smiles_list is not None: all_info_dict["smiles_list"] = input_smiles_list else: raise ValueError( "Either do_molecular_design must be True/False or input_smiles_list must be provided." ) # Retrosynthesis if do_retrosynthesis: if all_info_dict["smiles_list"] is None: raise ValueError( "Either molecular design must be performed or input_smiles_list must be provided for retrosynthesis." ) all_info_dict["retro_plan_dict"] = {} for i, smiles in enumerate(all_info_dict["smiles_list"]): if design_text_list is not None: design_text = design_text_list[0] else: design_text = None all_info_dict["retro_plan_dict"][smiles] = self.retrosynthesize( input_ids[i] if input_ids.dim() > 1 else input_ids, smiles, molecule_graphs=molecule_graphs, starting_mols=starting_mols, expansion_topk=expansion_topk, iterations=iterations, molecule_cost_weight=molecule_cost_weight, language_cost_weight=language_cost_weight, max_planning_time=max_planning_time, design_text=design_text, **kwargs, ) else: all_info_dict["retro_plan_dict"] = { smile: {"success": None} for smile in all_info_dict["smiles_list"] } for batch_idx, generated_mol in enumerate(all_info_dict["smiles_list"]): token_list = [] text_list = [] ignore_positions = {} if do_molecular_design: design_tokens = all_info_dict["design_analysis_tokens"][ batch_idx ].tolist() token_list = design_tokens + [IGNORE_INDEX] if generated_mol is None: generated_mol = "" text_list = [ self.tokenizer.decode( design_tokens, skip_special_tokens=True, clean_up_tokenization_spaced=True, ), generated_mol + ". ", ] ignore_positions = {0: generated_mol} if do_retrosynthesis: available_mols = self.graph_predictor.available["smiles"].tolist() retro_plan = all_info_dict["retro_plan_dict"][generated_mol] if retro_plan["success"] is not None and retro_plan["success"]: for i, (reaction, template, cost, analysis_tokens) in enumerate( zip( retro_plan["reaction_list"], retro_plan["templates"], retro_plan["cost"], retro_plan["analysis_tokens"], ) ): if isinstance(analysis_tokens, torch.Tensor): analysis_tokens = analysis_tokens.tolist() token_list.extend(analysis_tokens + [IGNORE_INDEX]) text_list.extend( [ self.tokenizer.decode( analysis_tokens, skip_special_tokens=True, clean_up_tokenization_spaced=True, ), reaction if reaction is not None else "", " with the template ", template if template is not None else "", " which requires the reactants: ", ] ) # Add these two lines to extract and add reactants if reaction is not None: reactants = reaction.split(">>")[1].split(".") formatted_reactants = [] for reactant in reactants: if reactant in available_mols: formatted_reactants.append( f"{reactant} (available)" ) else: formatted_reactants.append(reactant) text_list.extend([", ".join(formatted_reactants), ". "]) else: text_list.extend([". "]) ignore_positions[len(token_list) - 1] = ( reaction, template, cost, ) else: analysis_tokens = retro_plan["analysis_tokens"] if isinstance(analysis_tokens, torch.Tensor): analysis_tokens = analysis_tokens.tolist() token_list.extend(analysis_tokens) text_list.extend( [ self.tokenizer.decode( analysis_tokens, skip_special_tokens=True, clean_up_tokenization_spaced=True, ), " ", ] ) all_info_dict["token_lists"].append(token_list) all_info_dict["text_lists"].append(text_list) all_info_dict[f"batch_{batch_idx}_ignore_positions"] = ignore_positions all_info_dict["IGNORE_INDEX"] = IGNORE_INDEX return all_info_dict