| """Minimal padded-batch smoke test for the ApexOracle DLM encoder.""" | |
| import torch | |
| from transformers import AutoTokenizer | |
| from DLM_emb_model import MolEmbDLM | |
| MODEL_DIR = "Kiria-Nozan/ApexOracle" | |
| device = torch.device("cuda") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR) | |
| model = MolEmbDLM.from_pretrained(MODEL_DIR).eval().to(device) | |
| batch = tokenizer( | |
| ["[C] [C] [O]", "[C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1]"], | |
| padding=True, | |
| truncation=False, | |
| return_tensors="pt", | |
| ).to(device) | |
| with torch.no_grad(): | |
| hidden_states = model(**batch) | |
| print(hidden_states.shape) | |