Spaces:
Sleeping
Sleeping
| import warnings, os, datetime | |
| warnings.filterwarnings("ignore") | |
| import json | |
| import queue | |
| import threading | |
| import torch | |
| import numpy as np | |
| from tqdm import tqdm | |
| import pyarrow.parquet as pq | |
| from rdkit import Chem | |
| import torch.distributed as dist | |
| from torch.nn.parallel import DistributedDataParallel as DDP | |
| from model.CSU_EI import CSUEI | |
| from model.tokenizer import build_tokenizer | |
| from rdkit import RDLogger | |
| RDLogger.DisableLog('rdApp.*') | |
| def init_environment(): | |
| if "WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 1: | |
| dist.init_process_group(backend="nccl", timeout=datetime.timedelta(hours=2)) | |
| local_rank = int(os.environ["LOCAL_RANK"]) | |
| rank = dist.get_rank() | |
| world_size = dist.get_world_size() | |
| torch.cuda.set_device(local_rank) | |
| device = torch.device(f"cuda:{local_rank}") | |
| return True, rank, local_rank, world_size, device | |
| else: | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| return False, 0, 0, 1, device | |
| def distributed_extract_features(parquet_path, config_json_path, read_chunk_size=50000, per_gpu_batch=4096, precision='float16'): | |
| is_ddp, rank, local_rank, world_size, device = init_environment() | |
| if is_ddp and rank != 0: | |
| import sys | |
| sys.stdout = open(os.devnull, 'w') | |
| sys.stderr = open(os.devnull, 'w') | |
| exp_dir = os.path.dirname(config_json_path) | |
| time_stamp = '_'.join(exp_dir.split('_')[-2:]) | |
| with open(config_json_path, 'r') as f: | |
| config_json = json.load(f) | |
| config_json["model_config"]["roberta_model_path"] = None | |
| model_path = os.path.join(exp_dir, f'checkpoints/trained_model_{time_stamp}.pth') | |
| model_dict = torch.load(model_path, map_location=device) | |
| model = CSUEI(**config_json["model_config"]) | |
| model.load_state_dict(model_dict, strict=False) | |
| model.to(device) | |
| if is_ddp: | |
| model = DDP(model, device_ids=[local_rank]) | |
| smiles_encoder = model.module.smiles_encoder | |
| else: | |
| smiles_encoder = model.smiles_encoder | |
| model.eval() | |
| tokenizer = build_tokenizer(config_json['roberta_tokenizer_path']) | |
| max_len = config_json["dataset_config"]["smiles_maxlen"] | |
| if rank == 0: | |
| print(f"Initializing Pipelined CPU-GPU Stream processing from {parquet_path}...") | |
| pf = pq.ParquetFile(parquet_path) | |
| total_len = pf.metadata.num_rows | |
| feature_dim = config_json["model_config"]["feature_dim"] | |
| parent_folder = os.path.basename(os.path.dirname(parquet_path)) | |
| output_dir = f"database/molecular_vector_lib/{parent_folder}_{time_stamp}" | |
| os.makedirs(output_dir, exist_ok=True) | |
| npy_path = os.path.join(output_dir, f"{parent_folder}_DB_smiles_feature_{precision}_{time_stamp}.bin") | |
| if rank == 0: | |
| if os.path.exists(npy_path): | |
| os.remove(npy_path) | |
| print(f"Pre-allocating {precision} memory-mapped file on disk ({total_len} rows)...") | |
| dummy = np.memmap(npy_path, dtype=precision, mode='w+', shape=(total_len, feature_dim)) | |
| dummy.flush() | |
| del dummy | |
| if is_ddp: | |
| dist.barrier() | |
| features_mmap = np.memmap(npy_path, dtype=precision, mode='r+', shape=(total_len, feature_dim)) | |
| if rank == 0: | |
| print(f"Starting pipelined streaming (Disk Chunk Size: {read_chunk_size}, GPU Batch Size: {per_gpu_batch})...") | |
| data_queue = queue.Queue(maxsize=4) | |
| def cpu_producer(): | |
| try: | |
| global_idx = 0 | |
| batch_ids, batch_masks, batch_idxs = [], [], [] | |
| for batch in pf.iter_batches(batch_size=read_chunk_size, columns=['smiles']): | |
| smiles_col = batch.column('smiles').to_pylist() | |
| for smi in smiles_col: | |
| absolute_idx = global_idx | |
| global_idx += 1 | |
| if is_ddp and (absolute_idx % world_size != rank): | |
| continue | |
| if not isinstance(smi, str): | |
| continue | |
| mol = Chem.MolFromSmiles(smi) | |
| if mol is None: | |
| continue | |
| encode_dict = tokenizer( | |
| Chem.MolToSmiles(mol), | |
| max_length=max_len, | |
| padding='max_length', | |
| truncation=True | |
| ) | |
| batch_ids.append(encode_dict['input_ids']) | |
| batch_masks.append(encode_dict['attention_mask']) | |
| batch_idxs.append(absolute_idx) | |
| if len(batch_ids) == per_gpu_batch: | |
| data_queue.put((batch_ids, batch_masks, batch_idxs)) | |
| batch_ids, batch_masks, batch_idxs = [], [], [] | |
| if batch_ids: | |
| data_queue.put((batch_ids, batch_masks, batch_idxs)) | |
| except Exception as e: | |
| print(f"[Rank {rank}] CPU Producer error: {e}") | |
| finally: | |
| data_queue.put(None) | |
| producer_thread = threading.Thread(target=cpu_producer) | |
| producer_thread.daemon = True | |
| producer_thread.start() | |
| with torch.no_grad(): | |
| pbar = tqdm(total=total_len, desc=f"Pipeline Extracting (Rank {rank})", disable=(rank != 0)) | |
| while True: | |
| item = data_queue.get() | |
| if item is None: | |
| break | |
| batch_ids, batch_masks, batch_idxs = item | |
| current_batch_size = len(batch_idxs) | |
| _process_and_write(smiles_encoder, device, batch_ids, batch_masks, batch_idxs, features_mmap, precision) | |
| pbar.update(current_batch_size * world_size if is_ddp else current_batch_size) | |
| pbar.close() | |
| producer_thread.join() | |
| if rank == 0: | |
| print("Flushing final data to disk...") | |
| features_mmap.flush() | |
| del features_mmap | |
| if is_ddp: | |
| dist.barrier() | |
| dist.destroy_process_group() | |
| if rank == 0: | |
| print(f"Inference completed. {precision.upper()} binary file saved to {npy_path}") | |
| print(f"Done! All assets saved in {output_dir}") | |
| def _process_and_write(smiles_encoder, device, input_ids_list, attn_mask_list, indices, features_mmap, precision): | |
| input_ids = torch.tensor(input_ids_list, device=device) | |
| attn_mask = torch.tensor(attn_mask_list, device=device) | |
| with torch.amp.autocast('cuda'): | |
| feat, _ = smiles_encoder([input_ids, attn_mask]) | |
| feat = feat / feat.norm(dim=-1, keepdim=True) | |
| features_mmap[np.array(indices)] = feat.cpu().numpy().astype(precision) | |
| if __name__ == "__main__": | |
| parquet_path = "database/molecular_dataset/CSU-EI-DB_100M/id2molecule.parquet" | |
| config_path = "results/CSUEI_nist23_smil6_msl6_bs1024_DDP1_lr4.0e-04_smiesa1_msesa1_rm1_se1_2026-02-13_193603/config_2026-02-13_193603.json" | |
| distributed_extract_features( | |
| parquet_path, | |
| config_path, | |
| read_chunk_size=500000, | |
| per_gpu_batch=4096, | |
| precision='float16' | |
| ) | |