solve same embedding bug

DLM_emb_model.py

@@ -31,10 +31,10 @@ import ast
 from omegaconf import OmegaConf, DictConfig, ListConfig
 from huggingface_hub import PyTorchModelHubMixin
 
-# current_directory = Path(__file__).parent
-current_directory = Path('/data2/tianang/projects/Synergy')
+current_directory = Path(__file__).parent
+# current_directory = Path('/data2/tianang/projects/Synergy')
 
-with initialize_config_dir(config_dir="/data2/tianang/projects/mdlm/configs"):
+with initialize_config_dir(config_dir=str(current_directory/"configs")):
     config = compose(config_name="config")
 
 class mol_emb_mdlm(nn.Module):

compare_source_vs_hf.py

@@ -0,0 +1,108 @@
+
+import torch
+from transformers import AutoTokenizer
+import sys
+import os
+from hydra import compose, initialize_config_dir
+from pathlib import Path
+import numpy as np
+
+# Add current dir to path
+sys.path.append(os.getcwd())
+
+try:
+    from DLM_emb_model import MolEmbDLM
+except ImportError:
+    print("Could not import MolEmbDLM. Make sure you are running from ApexOracle directory.")
+    exit(1)
+
+def load_source_model():
+    print("Loading Source Model...")
+    current_directory = Path(os.getcwd())
+    # Replicating logic from DLM_emb_model.py
+    with initialize_config_dir(config_dir=str(current_directory/"configs"), version_base=None):
+        config = compose(config_name="config")
+    
+    model_name = "ibm-research/materials.selfies-ted"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    
+    DIT_ckpt_path = '/data2/tianang/projects/mdlm/Checkpoints_fangping/1-255000-fine-tune.ckpt'
+    model = MolEmbDLM(config, len(tokenizer.get_vocab()), DIT_ckpt_path, tokenizer.mask_token_id)
+    model.eval()
+    return model, tokenizer
+
+def load_hf_model():
+    print("Loading HF Model...")
+    model_path = "/data2/tianang/projects/mdlm/huggingface/huggingface_model"
+    # We use the same class but loaded via from_pretrained
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        model = MolEmbDLM.from_pretrained(model_path)
+    except Exception as e:
+        print(f"Failed to load HF model: {e}")
+        # Fallback to local if needed, though path is absolute
+        model = MolEmbDLM.from_pretrained(".")
+    model.eval()
+    return model, tokenizer
+
+def main():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    
+    # Load Source Model
+    source_model, source_tokenizer = load_source_model()
+    source_model.to(device)
+    
+    # Load HF Model
+    hf_model, hf_tokenizer = load_hf_model()
+    hf_model.to(device)
+    
+    # Test Input (SELFIES)
+    selfies = "[C][C][=O][O]" # Ethanol  "[C][C][=O][O]"
+    processed_selfies = selfies.replace('][', '] [')
+    
+    print(f"Testing with SELFIES: {processed_selfies}")
+    
+    # Tokenize (using source tokenizer for both to ensure identical input ids if tokenizers are same)
+    # Note: HF model folder has its own tokenizer files, source uses "ibm-research/materials.selfies-ted".
+    # They should be the same, but let's verify input_ids match too.
+    
+    inputs_source = source_tokenizer(processed_selfies, return_tensors="pt", padding=False, truncation=False)
+    inputs_hf = hf_tokenizer(processed_selfies, return_tensors="pt", padding=False, truncation=False)
+    
+    print(f"Source Input IDs: {inputs_source['input_ids']}")
+    print(f"HF Input IDs:     {inputs_hf['input_ids']}")
+    
+    if not torch.equal(inputs_source['input_ids'], inputs_hf['input_ids']):
+        print("WARNING: Tokenizers produced different input IDs!")
+    
+    # Run Source Model
+    inputs_s = {k: v.to(device) for k, v in inputs_source.items() if k in ["input_ids", "attention_mask"]}
+    with torch.no_grad():
+        emb_source = source_model(**inputs_s)
+        
+    # Run HF Model
+    inputs_h = {k: v.to(device) for k, v in inputs_hf.items() if k in ["input_ids", "attention_mask"]}
+    with torch.no_grad():
+        emb_hf = hf_model(**inputs_h)
+
+    print(f'Huggingface Embeddings: {emb_hf[0][0]}')
+        
+    print(f"Source Emb Shape: {emb_source.shape}")
+    print(f"HF Emb Shape:     {emb_hf.shape}")
+    
+    # Compare
+    diff = torch.abs(emb_source - emb_hf).sum().item()
+    max_diff = torch.abs(emb_source - emb_hf).max().item()
+    
+    print(f"Sum of Absolute Differences: {diff}")
+    print(f"Max Absolute Difference:     {max_diff}")
+    
+    if diff < 1e-5: # Allow small floating point differences
+        print("SUCCESS: Embeddings are identical (or extremely close).")
+    else:
+        print("FAILURE: Embeddings differ significantly.")
+        print(f"Source Mean: {emb_source.mean().item()}")
+        print(f"HF Mean:     {emb_hf.mean().item()}")
+
+if __name__ == "__main__":
+    main()

configs/callbacks/checkpoint_every_n_steps.yaml

@@ -0,0 +1,8 @@
+checkpoint_every_n_steps:
+  _target_: lightning.pytorch.callbacks.ModelCheckpoint
+  save_top_k: -1 # Do not save any "best" models; this callback is being used to save every n train steps
+  save_last: True # save model as ${save_dir}/checkpoints/last.ckpt
+  dirpath: ${checkpointing.save_dir}/checkpoints
+  verbose: True
+  auto_insert_metric_name: False
+  every_n_train_steps: 500

configs/callbacks/checkpoint_monitor.yaml

@@ -0,0 +1,10 @@
+checkpoint_monitor:
+  _target_: lightning.pytorch.callbacks.ModelCheckpoint
+  monitor: val/nll # name of the logged metric which determines when model is improving
+  mode: min # can be "max" or "min"
+  save_top_k: 1 # save k best models (determined by above metric)
+  save_last: False # True = additionally always save model from last epoch
+  dirpath: ${checkpointing.save_dir}/checkpoints
+  filename: best
+  auto_insert_metric_name: False
+  verbose: True

configs/callbacks/learning_rate_monitor.yaml

@@ -0,0 +1,3 @@
+learning_rate_monitor:
+  _target_: lightning.pytorch.callbacks.LearningRateMonitor
+  logging_interval: step

configs/config.yaml

@@ -0,0 +1,102 @@
+defaults:
+  - _self_
+  - /callbacks: [checkpoint_every_n_steps, checkpoint_monitor, learning_rate_monitor]
+  - /data: openwebtext
+  - /model: small  # small / medium
+  - /strategy: ddp
+  - /noise: loglinear
+  - /lr_scheduler: constant_warmup
+
+mode: sample_eval  # train / ppl_eval / sample_eval
+diffusion: absorbing_state
+backbone: dit  # dit / dimamba / ar
+parameterization: subs  # subs / d3pm / sedd
+time_conditioning: False
+T: 0  # 0 (continuous time) / 1000 
+subs_masking: False
+
+seed: 1
+
+loader:
+  global_batch_size: 512
+  eval_global_batch_size: ${.global_batch_size}
+  # Note: batch_size and eval_batch_size are **per machine**
+  batch_size: ${div_up:${.global_batch_size}, ${eval:${trainer.devices} * ${trainer.num_nodes}}}
+  eval_batch_size: ${div_up:${.eval_global_batch_size}, ${eval:${trainer.devices} * ${trainer.num_nodes}}}
+  num_workers: ${eval:"len(__import__('os').sched_getaffinity(0))"}
+  pin_memory: True
+
+sampling:
+  predictor: ddpm_cache  # analytic, ddpm, ddpm_cache
+  steps: 128
+  noise_removal: True
+  # TODO(yair): @subham, why aren't these params under `eval`?
+  num_sample_batches: 2  # Total samples: `num_gpus` * `loader.eval_batch_size` * num_sample_batches
+  num_sample_log: 2
+  semi_ar: False
+  stride_length: 1
+  num_strides: 1
+
+
+training:
+  ema: 0.9999
+  antithetic_sampling: True
+  importance_sampling: False
+  sampling_eps: 1e-3
+  change_of_variables: False
+
+eval:
+  checkpoint_path: '/data2/tianang/projects/mdlm/Checkpoints_fangping/1-255000-fine-tune.ckpt'  # Used to evaluate a checkpoint after training.
+  disable_ema: False
+  compute_generative_perplexity: False
+  perplexity_batch_size: 8
+  compute_perplexity_on_sanity: False
+  gen_ppl_eval_model_name_or_path: gpt2-large  # gpt2-large, meta-llama/Llama-2-7b-hf
+  generate_samples: True
+
+optim:
+  weight_decay: 0
+  lr: 3e-4
+  beta1: 0.9
+  beta2: 0.999
+  eps: 1e-8
+
+trainer:
+  _target_: lightning.Trainer
+  accelerator: cuda
+  num_nodes: 1
+  devices: ${device_count:}
+  accumulate_grad_batches: ${div_up:${loader.global_batch_size}, ${eval:${trainer.devices} * ${loader.batch_size} * ${trainer.num_nodes}}}
+  gradient_clip_val: 1.0
+  precision: 'bf16'
+  num_sanity_val_steps: 2
+  max_steps: 1_000_000
+  log_every_n_steps: 10
+  limit_train_batches: 1.0   # train on full dataset, can be used to toggle quick run
+  limit_val_batches: 1.0     # validate on full dataset, can be used to toggle quick run
+  val_check_interval: 10000
+
+wandb:
+  project: text-diffusion
+  notes: Mulan for text
+  group: null
+  job_type: null
+  name: null
+  id: ${.name}_${seed}
+  tags:
+    - ${noise.type}
+    - ${data.train}
+    - ${data.valid}
+
+hydra:
+  run:
+    dir: ./outputs/${data.train}/${now:%Y.%m.%d}/${now:%H%M%S}
+  job:
+    chdir: true
+
+checkpointing:
+  # Use custom `save_dir` if, e.g., saving to S3 bucket, otherwise leave this parameter as is
+  save_dir: ${cwd:}
+  # Note: `checkpoints` path should correspond to `checkpoint_every_n_steps.dirpath`
+  resume_from_ckpt: true
+  resume_ckpt_path: ${.save_dir}/checkpoints/last.ckpt

configs/data/ag_news.yaml

@@ -0,0 +1,6 @@
+train: ag_news
+valid: ag_news
+tokenizer_name_or_path: gpt2
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/lambada.yaml

@@ -0,0 +1,6 @@
+train: lambada
+valid: lambada
+tokenizer_name_or_path: gpt2
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/lm1b-gpt2.yaml

@@ -0,0 +1,6 @@
+train: lm1b
+valid: lm1b
+tokenizer_name_or_path: gpt2
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/lm1b-streaming.yaml

@@ -0,0 +1,6 @@
+train: lm1b
+valid: lm1b
+tokenizer_name_or_path: bert-base-uncased
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: False
+streaming: True

configs/data/lm1b.yaml

@@ -0,0 +1,6 @@
+train: lm1b
+valid: lm1b
+tokenizer_name_or_path: bert-base-uncased
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: False
+streaming: False

configs/data/openwebtext-split.yaml

@@ -0,0 +1,6 @@
+train: openwebtext-train
+valid: openwebtext-valid
+tokenizer_name_or_path: gpt2
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/openwebtext-streaming.yaml

@@ -0,0 +1,6 @@
+train: openwebtext
+valid: wikitext103
+tokenizer_name_or_path: gpt2
+cache_dir: /tmp/data
+wrap: True
+streaming: True

configs/data/openwebtext.yaml

@@ -0,0 +1,6 @@
+train: openwebtext
+valid: wikitext103
+tokenizer_name_or_path: ibm-research/materials.selfies-ted
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/ptb.yaml

@@ -0,0 +1,6 @@
+train: ptb
+valid: ptb
+tokenizer_name_or_path: gpt2
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/scientific_papers_arxiv.yaml

@@ -0,0 +1,6 @@
+train: scientific_papers_arxiv
+valid: scientific_papers_arxiv
+tokenizer_name_or_path: gpt2
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/scientific_papers_pubmed.yaml

@@ -0,0 +1,6 @@
+train: scientific_papers_pubmed
+valid: scientific_papers_pubmed
+tokenizer_name_or_path: gpt2
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/text8-crop.yaml

@@ -0,0 +1,7 @@
+# TODO: When using this dataset, set model.length = 256 to match D3PM setup
+train: text8-crop
+valid: text8
+tokenizer_name_or_path: text8
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/text8.yaml

@@ -0,0 +1,7 @@
+# TODO: When using this dataset, set model.length = 256 to match D3PM setup
+train: text8
+valid: text8
+tokenizer_name_or_path: text8
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/wikitext103.yaml

@@ -0,0 +1,6 @@
+train: wikitext103
+valid: wikitext103
+tokenizer_name_or_path: gpt2
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/data/wikitext2.yaml

@@ -0,0 +1,6 @@
+train: wikitext2
+valid: wikitext2
+tokenizer_name_or_path: gpt2
+cache_dir: /share/kuleshov/ssahoo/textdiffusion/data
+wrap: True
+streaming: False

configs/lr_scheduler/constant_warmup.yaml

@@ -0,0 +1,2 @@
+_target_: transformers.get_constant_schedule_with_warmup
+num_warmup_steps: 2500

configs/lr_scheduler/cosine_decay_warmup.yaml

@@ -0,0 +1,7 @@
+_target_: utils.CosineDecayWarmupLRScheduler
+t_in_epochs: False
+t_initial: ${eval:${trainer.max_steps}-${.warmup_t}}
+warmup_prefix: True
+warmup_lr_init: 1e-6
+warmup_t: ${eval:0.1*${trainer.max_steps}}
+lr_min: 1e-6

configs/model/medium.yaml

@@ -0,0 +1,10 @@
+name: medium
+type: ddit
+hidden_size: 1024
+cond_dim: 128
+length: 1024
+n_blocks: 24
+n_heads: 16
+scale_by_sigma: True
+dropout: 0.1
+tie_word_embeddings: False

configs/model/small-ar.yaml

@@ -0,0 +1,11 @@
+name: small
+type: ddit
+hidden_size: 768
+cond_dim: 128
+length: 1024
+n_blocks: 12
+n_heads: 12
+scale_by_sigma: True
+dropout: 0.1
+causal: True
+tie_word_embeddings: False

configs/model/small.yaml

@@ -0,0 +1,10 @@
+name: small
+type: ddit
+hidden_size: 768
+cond_dim: 128
+length: 1024
+n_blocks: 12
+n_heads: 12
+scale_by_sigma: True
+dropout: 0.1
+tie_word_embeddings: False

configs/model/tiny-ar.yaml

@@ -0,0 +1,11 @@
+name: tiny
+type: ddit
+hidden_size: 512
+cond_dim: 128
+length: 1024
+n_blocks: 8
+n_heads: 8
+scale_by_sigma: True
+dropout: 0.1
+causal: True
+tie_word_embeddings: False

configs/model/tiny-dimamba.yaml

@@ -0,0 +1,11 @@
+name: tiny
+type: dimamba
+hidden_size: 512
+cond_dim: 128
+length: 1024
+n_blocks: 14
+n_heads: 8
+scale_by_sigma: True
+dropout: 0.1
+temb_strategy: adaln 
+tie_word_embeddings: False

configs/model/tiny.yaml

@@ -0,0 +1,10 @@
+name: tiny
+type: ddit
+hidden_size: 512
+cond_dim: 128
+length: 1024
+n_blocks: 8
+n_heads: 8
+scale_by_sigma: True
+dropout: 0.1
+tie_word_embeddings: False

configs/noise/ar.yaml

@@ -0,0 +1,2 @@
+type: ar
+scale: 6.0

configs/noise/linear.yaml

@@ -0,0 +1,3 @@
+type: linear
+sigma_min: 1e-3
+sigma_max: 7.0

configs/noise/loglinear.yaml

@@ -0,0 +1,3 @@
+type: loglinear
+sigma_min: 1e-4
+sigma_max: 20

configs/noise/polynomial.yaml

@@ -0,0 +1,5 @@
+type: polynomial
+a: -3
+b: 5
+c: -4
+eps: 1e-3

configs/strategy/ddp.yaml

@@ -0,0 +1,2 @@
+_target_: lightning.pytorch.strategies.DDPStrategy
+find_unused_parameters: false  # TODO(yair): this seems hacky, I think if things are correct we shouldn't need this

configs/strategy/fsdp.yaml

@@ -0,0 +1,3 @@
+# TODO(yair): Currenly not compatible with grad clipping
+_target_: lightning.pytorch.strategies.FSDPStrategy
+sharding_strategy: SHARD_GRAD_OP

models/dit.py

@@ -339,7 +339,7 @@ class DDiTBlock_non_pad(nn.Module):
     qkv = rearrange(qkv, 'b s ... -> (b s) ...')
 
     # --------------------------------
-    mask_flat = attnmask.reshape(-1)
+    mask_flat = attnmask.reshape(-1).bool()
     qkv = qkv[mask_flat]
     seqlens = attnmask.sum(dim=1)
     pad_seq_len = torch.zeros(len(seqlens)+1, dtype=torch.int32, device=qkv.device)

reproduce_issue.py

@@ -0,0 +1,71 @@
+
+import torch
+from transformers import AutoTokenizer
+import sys
+import os
+
+# Add current dir to path
+sys.path.append(os.getcwd())
+
+try:
+    from DLM_emb_model import MolEmbDLM
+except ImportError:
+    print("Could not import MolEmbDLM. Make sure you are running from ApexOracle directory.")
+    exit(1)
+
+model_path = "/data2/tianang/projects/mdlm/huggingface/huggingface_model"
+
+print(f"Loading model from {model_path}...")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = MolEmbDLM.from_pretrained(model_path)
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    # Try loading from local directory if remote fails (though path is local)
+    try:
+        print("Trying to load from local directory...")
+        model = MolEmbDLM.from_pretrained(".")
+    except Exception as e2:
+        print(f"Failed to load from local: {e2}")
+        exit(1)
+
+model.eval()
+model.to("cuda" if torch.cuda.is_available() else "cpu")
+
+# Two different molecules
+smiles_list = [
+    "CC(=O)OC1=CC=CC=C1C(=O)O", # Aspirin
+    "CN1C=NC2=C1C(=O)N(C(=O)N2C)C" # Caffeine
+]
+
+print("Tokenizing inputs...")
+inputs = tokenizer(smiles_list, return_tensors="pt", padding=True, truncation=True)
+inputs = {k: v.to(next(model.parameters()).device) for k, v in inputs.items() if k in ["input_ids", "attention_mask"]}
+# Force different inputs to verify model logic, bypassing tokenizer issue
+inputs['input_ids'][1] = inputs['input_ids'][1] + 1
+print(f"Input IDs: {inputs['input_ids']}")
+print(f"Attention Mask: {inputs['attention_mask']}")
+
+print("Running model...")
+with torch.no_grad():
+    outputs = model(**inputs)
+
+# outputs is the tensor returned by forward
+embeddings = outputs
+
+print(f"Embeddings shape: {embeddings.shape}")
+
+emb1 = embeddings[0].cpu().numpy()
+emb2 = embeddings[1].cpu().numpy()
+
+# Compare
+diff = torch.abs(embeddings[0] - embeddings[1]).sum().item()
+print(f"Difference between embeddings: {diff}")
+
+if diff < 1e-6:
+    print("ISSUE REPRODUCED: Embeddings are identical.")
+else:
+    print("Embeddings are different.")
+
+print(f"Emb1 mean: {emb1.mean()}")
+print(f"Emb2 mean: {emb2.mean()}")

temp_data/polymers_lit_scraped.csv

@@ -0,0 +1,57 @@
+Notebook reference,Polymer name,monomer A,mol fraction A,monomer B,fraction B,monomer C,fraction C,monomer D,fraction D,monomer E,fraction E,monomer F,fraction F,Distribution,Architecture,Target DP,MIC (E. coli),MIC (S. aureus),MIC (K. pneumoniae),MIC (E. faecium),HC50
+SW1.84.1,L-Ni31Mo10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.45,C=CC(=O)NC(C)C,0.43,C=CC(=O)N1CCOCC1,0.12,,,,,,,statistical,linear,70,>512,>512,,,>2000
+SW1.84.2,L-Ni31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.45,C=CC(=O)NC(C)C,0.43,C=CC(=O)NCCCOC,0.11,,,,,,,statistical,linear,70,>512,>512,,,>2000
+SW1.84.3,L-Phe31Mo10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.50,C=CC(=O)Nc1ccccc1,0.37,C=CC(=O)N1CCOCC1,0.13,,,,,,,statistical,linear,70,256,>512,,,>2000
+SW1.89.1,L-Phe31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.51,C=CC(=O)Nc1ccccc1,0.37,C=CC(=O)NCCCOC,0.13,,,,,,,statistical,linear,70,256,>512,,,>2000
+SW1.89.2,L-Do31Mo10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.59,C=CC(=O)NCCCCCCCCCCCC,0.26,C=CC(=O)N1CCOCC1,0.15,,,,,,,statistical,linear,70,128,32-64,256,512,>2000
+SW1.89.3,L-Do31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.59,C=CC(=O)NCCCCCCCCCCCC,0.26,C=CC(=O)NCCCOC,0.15,,,,,,,statistical,linear,70,128,32,512,512,>2000
+SW1.110.1,L-Ni13Mo4,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.73,C=CC(=O)NC(C)C,0.21,C=CC(=O)N1CCOCC1,0.06,,,,,,,statistical,linear,70,>512,32,,,>2000
+SW1.110.2,L-Ni13Mep4,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.73,C=CC(=O)NC(C)C,0.21,C=CC(=O)NCCCOC,0.06,,,,,,,statistical,linear,70,>512,64-128,,,>2000
+SW1.110.3,L-Phe13Mo4,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.77,C=CC(=O)Nc1ccccc1,0.17,C=CC(=O)N1CCOCC1,0.06,,,,,,,statistical,linear,70,>512,32,,,>2000
+SW1.115.1,L-Phe13Mep4,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.77,C=CC(=O)Nc1ccccc1,0.17,C=CC(=O)NCCCOC,0.06,,,,,,,statistical,linear,70,>512,32,,,>2000
+SW1.115.2,L-Do13Mo4,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.83,C=CC(=O)NCCCCCCCCCCCC,0.11,C=CC(=O)N1CCOCC1,0.06,,,,,,,statistical,linear,70,256-512,32,,,<50
+SW1.115.3,L-Do13Mep4,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.83,C=CC(=O)NCCCCCCCCCCCC,0.11,C=CC(=O)NCCCOC,0.06,,,,,,,statistical,linear,70,256,32,256,256,>2000
+SW1.119.1,H-Ni31Mo10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.45,C=CC(=O)NC(C)C,0.43,C=CC(=O)N1CCOCC1,0.12,,,,,,,statistical,linear,115,>512,128,,,>8000
+SW1.119.2,H-Ni31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.45,C=CC(=O)NC(C)C,0.43,C=CC(=O)NCCCOC,0.11,,,,,,,statistical,linear,115,>512,>512,,,>8000
+SW1.119.3,H-Phe31Mo10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.50,C=CC(=O)Nc1ccccc1,0.37,C=CC(=O)N1CCOCC1,0.13,,,,,,,statistical,linear,115,256-512,128-256,64,>512,>8000
+SW1.125.1,H-Phe31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.51,C=CC(=O)Nc1ccccc1,0.37,C=CC(=O)NCCCOC,0.13,,,,,,,statistical,linear,115,256,>512,nd,,>8000
+SW1.119.5,H-Do31Mo10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.59,C=CC(=O)NCCCCCCCCCCCC,0.26,C=CC(=O)N1CCOCC1,0.15,,,,,,,statistical,linear,115,128,32,128-256,256,>8000
+SW1.119.6,H-Do31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.59,C=CC(=O)NCCCCCCCCCCCC,0.26,C=CC(=O)NCCCOC,0.15,,,,,,,statistical,linear,115,128,32,256,>512,6300
+SW2.3.1,L-Bam31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.48,C=CC(=O)NCCCC,0.40,C=CC(=O)NCCCOC,0.12,,,,,,,statistical,linear,70,>512,>512,,,>8000
+SW2.3.2,L-Bmam31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.52,C=CC(=O)NCOCCCC,0.35,C=CC(=O)NCCCOC,0.13,,,,,,,statistical,linear,70,256,>512,,,6200
+SW2.3.3,L-Tmb31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.54,C=CC(=O)NC(C)(C)CC(C)(C)C,0.32,C=CC(=O)NCCCOC,0.14,,,,,,,statistical,linear,70,64,64,,,<62.5
+SW2.3.4,L-Oct31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.54,C=CC(=O)NCCCCCCCC,0.32,C=CC(=O)NCCCOC,0.14,,,,,,,statistical,linear,70,256-128,64,256,>512,4700
+SW2.3.5,L-Olam31Mep10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.63,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.21,C=CC(=O)NCCCOC,0.16,,,,,,,statistical,linear,70,128,64-32,>512,>512,>8000
+SW3.56.1,L-Do30Mep5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.66,C=CC(=O)NCCCCCCCCCCCC,0.26,C=CC(=O)NCCCOC,0.07,,,,,,,statistical,linear,70,512,128,,,3400
+SW3.56.2,L-Tmb5Mo90,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.04,C=CC(=O)NC(C)(C)CC(C)(C)C,0.04,C=CC(=O)N1CCOCC1,0.93,,,,,,,statistical,linear,70,>512,>512,,,>4000
+SW3.56.3,L-Oct5Mep5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.87,C=CC(=O)NCCCCCCCC,0.05,C=CC(=O)NCCCOC,0.07,,,,,,,statistical,linear,70,>512,>512,,,>4000
+SW3.56.4,L-Phe15Mo30,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.46,C=CC(=O)Nc1ccccc1,0.18,C=CC(=O)N1CCOCC1,0.37,,,,,,,statistical,linear,70,>512,16,,,>4000
+SW4.14.2,L-Aeg5Phe25Mo50Mep20,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.038,C=CC(=O)Nc1ccccc1,0.246,C=CC(=O)N1CCOCC1,0.514,C=CC(=O)NCCCOC,0.203,,,,,statistical,linear,70,>512,>512,,,2200
+SW4.29.1,L-Do5Mo40Mep5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.416,C=CC(=O)NCCCCCCCCCCCC,0.036,C=CC(=O)N1CCOCC1,0.488,C=CC(=O)NCCCOC,0.060,,,,,statistical,linear,70,>512,>512,,,>4000
+SW4.29.2,L-Phe20Olam5Mep5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.645,C=CC(=O)Nc1ccccc1,0.259,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.030,C=CC(=O)NCCCOC,0.067,,,,,statistical,linear,70,128,32,,,>4000
+SW5.20.1,L-Do25,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.777,C=CC(=O)NCCCCCCCCCCCC,0.223,,,,,,,,,statistical,linear,70,64,,,,>4000
+SW5.20.2,L-Aeg10Olam30Mo60,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.091,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.164,C=CC(=O)N1CCOCC1,0.745,,,,,,,statistical,linear,70,>512,,,,>4000
+SW5.20.3,L-Ni25Phe20,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.427,C=CC(=O)NC(C)C,0.355,C=CC(=O)Nc1ccccc1,0.218,,,,,,,statistical,linear,70,>512,,,,>4000
+SW5.20.4,L-Bam40Oct5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.438,C=CC(=O)NCCCC,0.517,C=CC(=O)NCCCCCCCC,0.045,,,,,,,statistical,linear,70,32,,,,<500
+SW5.20.5,L-Phe23Oct5Mo55,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.126,C=CC(=O)Nc1ccccc1,0.239,C=CC(=O)N1CCOCC1,0.038,C=CC(=O)N1CCOCC1,0.597,,,,,statistical,linear,70,>512,,,,>4000
+SW5.24.1,L-Aeg10Phe20Olam25,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.450,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.107,C=CC(=O)Nc1ccccc1,0.281,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.161,,,,,statistical,linear,70,128,,,,1500
+SW5.24.2,L-Aeg20Ni35Tmb10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.266,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.163,C=CC(=O)NC(C)C,0.486,C=CC(=O)NC(C)(C)CC(C)(C)C,0.086,,,,,statistical,linear,70,64,,,,<500
+SW5.24.3,L-Phe35Olam10Mo20,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.292,C=CC(=O)Nc1ccccc1,0.410,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.054,C=CC(=O)N1CCOCC1,0.244,,,,,statistical,linear,70,128,,,,>4000
+SW5.24.4,L-Aeg17Tmb8Mo37,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.319,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.148,C=CC(=O)NC(C)(C)CC(C)(C)C,0.078,C=CC(=O)N1CCOCC1,0.455,,,,,statistical,linear,70,256,,,,<500
+SW5.24.5,L-Aeg20Ni20Olam25Mo5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.269,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.193,C=CC(=O)NC(C)C,0.328,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.144,C=CC(=O)N1CCOCC1,0.066,,,statistical,linear,70,256,,,,>4000
+SW5.41.1,L-Do10,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.912,C=CC(=O)NCCCCCCCCCCCC,0.088,,,,,,,,,statistical,linear,70,256,,,,>4000
+SW5.41.2,L-Phe15Do5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.759,C=CC(=O)Nc1ccccc1,0.200,C=CC(=O)NCCCCCCCCCCCC,0.041,,,,,,,statistical,linear,70,256,,,,>4000
+SW5.41.3,L-Aeg5Phe5Olam5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.845,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.053,C=CC(=O)Nc1ccccc1,0.070,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.032,,,,,statistical,linear,70,128,,,,>4000
+SW5.41.4,L-Ni20Do5Mep5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.593,C=CC(=O)NC(C)C,0.309,C=CC(=O)NCCCCCCCCCCCC,0.037,C=CC(=O)NCCCOC,0.061,,,,,statistical,linear,70,256,,,,>4000
+SW5.41.5,L-Phe20Olam5Mo15,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.530,C=CC(=O)Nc1ccccc1,0.248,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.028,C=CC(=O)N1CCOCC1,0.194,,,,,statistical,linear,70,128,,,,>4000
+SW5.42.1,L-Phe5Do5Mo50,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.321,C=CC(=O)Nc1ccccc1,0.056,C=CC(=O)NCCCCCCCCCCCC,0.035,C=CC(=O)N1CCOCC1,0.588,,,,,statistical,linear,70,>512,,,,>4000
+SW5.42.2,L-Aeg10Oct15Tmb5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.678,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.104,C=CC(=O)NCCCCCCCC,0.164,C=CC(=O)NC(C)(C)CC(C)(C)C,0.055,,,,,statistical,linear,70,128-256,,,,<500
+SW5.42.3,L-Do5Bam5Mo20Mep5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.570,C=CC(=O)NCCCCCCCCCCCC,0.038,C=CC(=O)NCCCC,0.071,C=CC(=O)N1CCOCC1,0.257,C=CC(=O)NCCCOC,0.063,,,statistical,linear,70,256,,,,>4000
+SW5.42.4,L-Aeg5Phe15Bam30Mo25,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.183,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.039,C=CC(=O)Nc1ccccc1,0.154,C=CC(=O)NCCCC,0.356,C=CC(=O)N1CCOCC1,0.268,,,statistical,linear,70,512,,,,>4000
+SW5.42.5,L-Phe5Olam10Bmam10Mep5,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.674,C=CC(=O)Nc1ccccc1,0.068,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.062,C=CC(=O)NCOCCCC,0.127,C=CC(=O)NCCCOC,0.070,,,statistical,linear,70,64,,,,>4000
+SW5.65.1,L-Aeg5Ni10Phe5Do30Mep15,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.309,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.047,C=CC(=O)NC(C)C,0.161,C=CC(=O)Nc1ccccc1,0.062,C=CC(=O)NCCCCCCCCCCCC,0.229,C=CC(=O)NCCCOC,0.191,statistical,linear,70,64,,,,3300
+SW5.65.5,L-Aeg10Ni15Bam10Olam20Mep20,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.206,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.089,C=CC(=O)NC(C)C,0.226,C=CC(=O)NCCCC,0.134,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.106,C=CC(=O)NCCCOC,0.238,statistical,linear,70,128,,,,1400
+SW5.65.7,L-Do15Bam15Oct10Mo30,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.245,C=CC(=O)NCCCCCCCCCCCC,0.106,C=CC(=O)NCCCC,0.199,C=CC(=O)NCCCCCCCC,0.092,C=CC(=O)N1CCOCC1,0.358,,,statistical,linear,70,128,,,,>4000
+SW5.65.8,L-Aeg10Ni5Do25Tmb10Mep35,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.122,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.088,C=CC(=O)NC(C)C,0.075,C=CC(=O)NCCCCCCCCCCCC,0.211,C=CC(=O)NC(C)(C)CC(C)(C)C,0.092,C=CC(=O)NCCCOC,0.412,statistical,linear,70,>512,,,,<500
+SW5.65.9,L-Ni10Do5Mo60,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.185,C=CC(=O)NC(C)C,0.135,C=CC(=O)NCCCCCCCCCCCC,0.032,C=CC(=O)N1CCOCC1,0.649,,,,,statistical,linear,70,>512,,,,>4000
+SW5.65.10,L-Aeg15Ni10Do10Olam10Mep35,C=CC(=O)NCC[N+](C)(C)C.[Cl-],0.167,C=CC(=O)NCCNC(N)=[NH2+].[Cl-],0.134,C=CC(=O)NC(C)C,0.152,C=CC(=O)NCCCCCCCCCCCC,0.072,C=CC(=O)NCCCCCCCC/C=C\CCCCCCCC,0.054,C=CC(=O)NCCCOC,0.421,statistical,linear,70,>512,,,,2500

temp_fangping.py

@@ -0,0 +1,127 @@
+import pandas as pd
+import numpy as np
+from DLM_emb_model import MolEmbDLM
+from transformers import AutoTokenizer
+import torch
+import selfies as sf
+
+MODEL_DIR = "Kiria-Nozan/ApexOracle"
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+model = MolEmbDLM.from_pretrained(MODEL_DIR)
+model.eval()
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# Load CSV data
+df = pd.read_csv("temp_data/polymers_lit_scraped.csv")
+
+# Extract all unique monomer SMILES
+monomer_columns = ["monomer A", "monomer B", "monomer C", "monomer D", "monomer E", "monomer F"]
+all_monomers = set()
+
+for col in monomer_columns:
+    if col in df.columns:
+        monomers = df[col].dropna().unique()
+        all_monomers.update(monomers)
+
+print(f"Total unique monomers: {len(all_monomers)}")
+
+# Convert SMILES to SELFIES and prepare for embedding
+monomer_selfies = {}
+valid_monomers = []
+
+for smiles in all_monomers:
+    try:
+        selfies = sf.encoder(smiles)
+        monomer_selfies[smiles] = selfies
+        valid_monomers.append((smiles, selfies))
+    except Exception as e:
+        print(f"Error converting {smiles} to SELFIES: {e}")
+
+print(f"Valid monomers for embedding: {len(valid_monomers)}")
+
+# Generate embeddings for all monomers
+monomer_embeddings = {}
+
+for smiles, selfies in valid_monomers:
+    # Prepare input similar to example.py
+    batch = tokenizer(
+        selfies.replace('][', '] ['),
+        padding="max_length",
+        max_length=1024,
+        truncation=True,
+        return_tensors="pt",
+    )
+    
+    batch = {k: v.to(device) for k, v in batch.items()}
+    
+    with torch.no_grad():
+        embeddings = model(
+            input_ids=batch["input_ids"],
+            attention_mask=batch["attention_mask"]+1-batch["attention_mask"],
+        )
+    
+    # Store the embedding (average pooling over sequence length)
+    monomer_embeddings[smiles] = embeddings[0][0].cpu().numpy()
+
+print(f"Generated embeddings for {len(monomer_embeddings)} monomers")
+print(f"Embedding shape: {list(monomer_embeddings.values())[0].shape}")
+
+# Check for identical embeddings
+print("\nChecking for identical embeddings...")
+embedding_list = list(monomer_embeddings.items())
+identical_pairs = []
+
+for i in range(len(embedding_list)):
+    for j in range(i + 1, len(embedding_list)):
+        smiles1, emb1 = embedding_list[i]
+        smiles2, emb2 = embedding_list[j]
+        
+        # Check if embeddings are identical (with small tolerance for floating point precision)
+        if np.allclose(emb1, emb2, rtol=1e-09, atol=1e-09):
+            identical_pairs.append((smiles1, smiles2))
+
+if identical_pairs:
+    print(f"Found {len(identical_pairs)} pairs of identical embeddings:")
+    for smiles1, smiles2 in identical_pairs:
+        print(f"  {smiles1} <-> {smiles2}")
+    
+    # Analyze the identical groups
+    print("\nAnalyzing identical embedding groups...")
+    
+    # Create groups of molecules with identical embeddings
+    identical_groups = {}
+    processed = set()
+    
+    for smiles1, smiles2 in identical_pairs:
+        if smiles1 not in processed and smiles2 not in processed:
+            # Find all molecules identical to smiles1
+            group = {smiles1, smiles2}
+            for other_smiles1, other_smiles2 in identical_pairs:
+                if other_smiles1 in group:
+                    group.add(other_smiles2)
+                elif other_smiles2 in group:
+                    group.add(other_smiles1)
+            
+            group_key = frozenset(group)
+            if group_key not in identical_groups:
+                identical_groups[group_key] = group
+                processed.update(group)
+    
+    print(f"Found {len(identical_groups)} groups of molecules with identical embeddings:")
+    for i, group in enumerate(identical_groups.values(), 1):
+        print(f"\nGroup {i} ({len(group)} molecules):")
+        for smiles in sorted(group):
+            selfies_repr = monomer_selfies.get(smiles, "N/A")
+            print(f"  SMILES: {smiles}")
+            print(f"  SELFIES: {selfies_repr}")
+            print()
+else:
+    print("No identical embeddings found.")
+
+# Save results
+np.save("temp_data/monomer_embeddings.npy", monomer_embeddings)
+print("Embeddings saved to monomer_embeddings.npy")

verify_selfies.py

@@ -0,0 +1,83 @@
+
+import torch
+from transformers import AutoTokenizer
+import sys
+import os
+
+# Add current dir to path to find DLM_emb_model
+sys.path.append(os.getcwd())
+
+try:
+    from DLM_emb_model import MolEmbDLM
+except ImportError:
+    print("Could not import MolEmbDLM. Make sure you are running from ApexOracle directory.")
+    exit(1)
+
+# Use local model path where we applied the fix
+model_path = "/data2/tianang/projects/mdlm/huggingface/huggingface_model"
+
+print(f"Loading model from {model_path}...")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = MolEmbDLM.from_pretrained(model_path)
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    # Try loading from local directory if remote fails
+    try:
+        model = MolEmbDLM.from_pretrained(".")
+    except Exception as e2:
+        print(f"Failed to load from local: {e2}")
+        exit(1)
+
+model.eval()
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# Two different SELFIES
+selfies_list = [
+    "[C][C][O]",          # Ethanol
+    "[C][C][=O][O]"       # Acetic Acid
+]
+
+# Preprocessing from example.py: seq.replace('][', '] [')
+processed_selfies = [s.replace('][', '] [') for s in selfies_list]
+
+print(f"Processed SELFIES: {processed_selfies}")
+
+# Tokenize with padding=True to create a batch (essential to test the bug fix)
+# example.py had padding=False because it was single sequence.
+print("Tokenizing inputs...")
+inputs = tokenizer(
+    processed_selfies,
+    padding=True,
+    truncation=True,
+    return_tensors="pt"
+)
+
+print(f"Input IDs:\n{inputs['input_ids']}")
+print(f"Attention Mask:\n{inputs['attention_mask']}")
+
+inputs = {k: v.to(device) for k, v in inputs.items() if k in ["input_ids", "attention_mask"]}
+
+print("Running model...")
+with torch.no_grad():
+    embeddings = model(**inputs)
+
+print(f"Embeddings shape: {embeddings.shape}")
+
+# Compare embeddings of the two molecules
+# We compare the mean embedding or the first token embedding
+emb1 = embeddings[0]
+emb2 = embeddings[1]
+
+# Calculate difference
+diff = torch.abs(emb1 - emb2).sum().item()
+print(f"Difference between embeddings (sum of abs diff): {diff}")
+
+if diff < 1e-6:
+    print("ISSUE: Embeddings are identical.")
+else:
+    print("SUCCESS: Embeddings are different.")
+
+print(f"Emb1 mean: {emb1.mean().item()}")
+print(f"Emb2 mean: {emb2.mean().item()}")