[Code scan] Use dictionary special-token ids for UniMol v1 atomic representations

[njzjz]

This issue is a result of a Codex global repository scan.

Summary

UniMol v1 atomic representation filtering hard-codes token ids 0, 1, and 2 as BOS/EOS/PAD. The Dictionary actually initializes [CLS]=0, [UNK]=1, [PAD]=2, [SEP]=3. As a result, [SEP] is kept as a fake atom and unknown atoms are dropped from atomic representations.

Code references

unimol_tools/unimol_tools/models/unimol.py

Lines 255 to 283 in 4596596

	if return_atomic_reprs:
	filtered_tensors = []
	filtered_coords = []
	for tokens, coord in zip(src_tokens, src_coord):
	filtered_tensor = tokens[
	(tokens != 0) & (tokens != 1) & (tokens != 2)
	] # filter out BOS(0), EOS(1), PAD(2)
	filtered_coord = coord[(tokens != 0) & (tokens != 1) & (tokens != 2)]
	filtered_tensors.append(filtered_tensor)
	filtered_coords.append(filtered_coord)
	lengths = [
	len(filtered_tensor) for filtered_tensor in filtered_tensors
	] # Compute the lengths of the filtered tensors
	cls_atomic_reprs = []
	atomic_symbols = []
	for i in range(len(all_repr)):
	atomic_reprs = encoder_rep[i, 1 : lengths[i] + 1, :]
	atomic_symbol = []
	for atomic_num in filtered_tensors[i]:
	atomic_symbol.append(self.dictionary.symbols[atomic_num])
	atomic_symbols.append(atomic_symbol)
	cls_atomic_reprs.append(atomic_reprs)
	return {
	'cls_repr': cls_repr,
	'atomic_symbol': atomic_symbols,
	'atomic_coords': filtered_coords,
	'atomic_reprs': cls_atomic_reprs,

unimol_tools/unimol_tools/data/dictionary.py

Lines 19 to 33 in 4596596

	bos="[CLS]",
	pad="[PAD]",
	eos="[SEP]",
	unk="[UNK]",
	extra_special_symbols=None,
	):
	self.bos_word, self.unk_word, self.pad_word, self.eos_word = bos, unk, pad, eos
	self.symbols = []
	self.count = []
	self.indices = {}
	self.specials = set()
	# initialize dictionary with special tokens
	for token in [bos, unk, pad, eos]:
	self.add_symbol(token, is_special=True)

Impact

return_atomic_reprs=True can return an extra [SEP] entry and omit legitimate unknown atoms, so atomic symbols, coordinates, and representations no longer align with the molecule atoms.

Suggested fix

Build the filter mask from self.dictionary.bos(), self.dictionary.eos(), and self.dictionary.pad(). Do not filter self.dictionary.unk(), because unknown atom tokens still correspond to atom positions.