Upload base files

ChemQ3MTP/FastChemTokenizerHF.py

@@ -0,0 +1,539 @@
+import torch
+import json
+import os
+from typing import List, Union, Optional, Tuple, Dict, Any
+from functools import lru_cache
+from collections.abc import Mapping
+
+
+# ------------------------------
+# BatchEncoding
+# ------------------------------
+class BatchEncoding(dict, Mapping):
+    """Minimal BatchEncoding compatible wrapper."""
+
+    def __init__(self, data: dict, tensor_type: Optional[str] = None):
+        data = {} if data is None else {k: v for k, v in data.items()}
+        super().__init__(data)
+        self.data = data
+        self.tensor_type = tensor_type
+        for k, v in data.items():
+            setattr(self, k, v)
+
+    def __getitem__(self, key): return self.data[key]
+    def __iter__(self): return iter(self.data)
+    def __len__(self): return len(self.data)
+    def keys(self): return self.data.keys()
+    def values(self): return self.data.values()
+    def items(self): return self.data.items()
+    def get(self, key, default=None): return self.data.get(key, default)
+
+    def to(self, device):
+        if self.tensor_type in ("pt", "torch"):
+            for k, v in list(self.data.items()):
+                if torch.is_tensor(v):
+                    self.data[k] = v.to(device)
+                    setattr(self, k, self.data[k])
+        return self
+
+    def cpu(self): return self.to("cpu")
+    def cuda(self): return self.to("cuda")
+    def detach(self):
+        if self.tensor_type in ("pt", "torch"):
+            for k, v in list(self.data.items()):
+                if torch.is_tensor(v):
+                    self.data[k] = v.detach()
+                    setattr(self, k, self.data[k])
+        return self
+
+    def __repr__(self):
+        keys = ", ".join(list(self.data.keys())[:10])
+        return f"BatchEncoding(keys=[{keys}], tensor_type={self.tensor_type})"
+
+
+# ------------------------------
+# Base class
+# ------------------------------
+class PreTrainedTokenizerBase:
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            if key.endswith('_token'):
+                setattr(self, f"_{key}", value)
+                setattr(self, f"{key}_id", None)
+        self.model_max_length = kwargs.get('model_max_length', 512)
+        self.padding_side = kwargs.get('padding_side', 'right')
+        self.truncation_side = kwargs.get('truncation_side', 'right')
+        self.chat_template = kwargs.get('chat_template')
+
+
+# ------------------------------
+# Trie node
+# ------------------------------
+class TrieNode:
+    __slots__ = ['children', 'token_id']
+    def __init__(self):
+        self.children = {}
+        self.token_id = None
+
+
+# ------------------------------
+# FastChemTokenizer
+# ------------------------------
+
+class FastChemTokenizer(PreTrainedTokenizerBase):
+    def __init__(self, token_to_id=None, vocab_file=None, **kwargs):
+        if vocab_file is not None:
+            with open(vocab_file, "r", encoding="utf-8") as f:
+                token_to_id = json.load(f)
+                token_to_id = {str(k): int(v) for k, v in token_to_id.items()}
+
+        self.token_to_id = token_to_id
+        self.id_to_token = {v: k for k, v in token_to_id.items()}
+
+        # Build trie
+        self.trie_root = self._build_trie(self.token_to_id)
+
+        # ✅ Call parent (sets token *strings*, may reset *_id to None)
+        super().__init__(
+            bos_token="<s>",
+            eos_token="</s>",
+            unk_token="<unk>",
+            pad_token="<pad>",
+            mask_token="<mask>",
+            model_max_length=kwargs.get("model_max_length", 512),
+            padding_side=kwargs.get("padding_side", "right"),
+            truncation_side=kwargs.get("truncation_side", "right"),
+            **kwargs,
+        )
+
+        # ✅ Re-map token strings → IDs from vocab
+        self.bos_token_id  = self.token_to_id.get("<s>", 0)
+        self.eos_token_id  = self.token_to_id.get("</s>", 1)
+        self.pad_token_id  = self.token_to_id.get("<pad>", 2)
+        self.unk_token_id  = self.token_to_id.get("<unk>", 3)
+        self.mask_token_id = self.token_to_id.get("<mask>", 4)
+
+        # Ensure reverse mapping always valid
+        self.id_to_token[self.bos_token_id]  = "<s>"
+        self.id_to_token[self.eos_token_id]  = "</s>"
+        self.id_to_token[self.pad_token_id]  = "<pad>"
+        self.id_to_token[self.unk_token_id]  = "<unk>"
+        self.id_to_token[self.mask_token_id] = "<mask>"
+
+        # Debug
+        print("✅ Special tokens bound:",
+              self.bos_token_id, self.eos_token_id, self.pad_token_id,
+              self.unk_token_id, self.mask_token_id)
+        
+        # ✅ Ensure token *strings* also exist (for decode fallback)
+        self.bos_token = "<s>"
+        self.eos_token = "</s>"
+        self.pad_token = "<pad>"
+        self.unk_token = "<unk>"
+        self.mask_token = "<mask>"
+
+
+    def _build_trie(self, token_to_id):
+        root = TrieNode()
+        for token, tid in token_to_id.items():
+            node = root
+            for char in token:
+                if char not in node.children:
+                    node.children[char] = TrieNode()
+                node = node.children[char]
+            node.token_id = tid
+        return root
+
+    @property
+    def vocab_size(self): return len(self.token_to_id)
+    def __len__(self): return len(self.token_to_id)
+    def get_vocab(self) -> Dict[str, int]: return self.token_to_id.copy()
+
+    @lru_cache(maxsize=10000)
+    def _cached_encode_str(self, s: str) -> Tuple[int, ...]:
+        return tuple(self._encode_core(s))
+
+    def _encode_core(self, text: str) -> List[int]:
+        tokens, result_ids = text, []
+        i, n = 0, len(tokens)
+        while i < n:
+            node, j = self.trie_root, i
+            last_match_id, last_match_end = None, i
+            while j < n and tokens[j] in node.children:
+                node = node.children[tokens[j]]
+                j += 1
+                if node.token_id is not None:
+                    last_match_id, last_match_end = node.token_id, j
+            if last_match_id is not None:
+                result_ids.append(last_match_id)
+                i = last_match_end
+            else:
+                tid = self.token_to_id.get(tokens[i], self.unk_token_id)
+                result_ids.append(tid)
+                i += 1
+        return result_ids
+
+    # ------------------------------
+    # Converters
+    # ------------------------------
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.token_to_id.get(token, self.unk_token_id)
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.id_to_token.get(index, self.unk_token)
+
+    def convert_tokens_to_ids(self, tokens: Union[str, List[str]]):
+        if isinstance(tokens, str): return self._convert_token_to_id(tokens)
+        return [self._convert_token_to_id(tok) for tok in tokens]
+
+    def convert_ids_to_tokens(self, ids: Union[int, List[int]]):
+        if isinstance(ids, int): return self._convert_id_to_token(ids)
+        return [self._convert_id_to_token(i) for i in ids]
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str: return "".join(tokens)
+
+    # ------------------------------
+    # Encoding / Decoding
+    # ------------------------------
+        # ------------------------------
+    # Convenience wrappers
+    # ------------------------------
+    def encode(
+        self,
+        text: str,
+        text_pair: Optional[str] = None,
+        add_special_tokens: bool = True,
+        padding: bool = False,
+        truncation: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+    ) -> List[int]:
+        encoded = self.encode_plus(
+            text=text,
+            text_pair=text_pair,
+            add_special_tokens=add_special_tokens,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+        )
+        input_ids = encoded["input_ids"]
+        if isinstance(input_ids, torch.Tensor):
+            if input_ids.dim() > 1:
+                input_ids = input_ids.squeeze(0)
+            input_ids = input_ids.tolist()
+        return input_ids
+
+    def __call__(
+        self, 
+        text: Union[str, List[str]], 
+        text_pair: Optional[Union[str, List[str]]] = None, 
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        is_split_into_words: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, Any]] = None,
+        return_token_type_ids: Optional[bool] = None,
+        return_attention_mask: Optional[bool] = None,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        **kwargs
+    ) -> BatchEncoding:
+        """HuggingFace-compatible: one string → encode_plus, list → batch_encode_plus"""
+        if return_token_type_ids is None:
+            return_token_type_ids = True
+        if return_attention_mask is None:
+            return_attention_mask = True
+
+        if isinstance(text, list):
+            if text_pair is not None:
+                batch = [(t, p) for t, p in zip(text, text_pair)]
+            else:
+                batch = text
+            return self.batch_encode_plus(
+                batch,
+                add_special_tokens=add_special_tokens,
+                padding=padding,
+                truncation=truncation,
+                max_length=max_length,
+                stride=stride,
+                is_split_into_words=is_split_into_words,
+                pad_to_multiple_of=pad_to_multiple_of,
+                return_tensors=return_tensors,
+                return_token_type_ids=return_token_type_ids,
+                return_attention_mask=return_attention_mask,
+                return_overflowing_tokens=return_overflowing_tokens,
+                return_special_tokens_mask=return_special_tokens_mask,
+                return_offsets_mapping=return_offsets_mapping,
+                return_length=return_length,
+                verbose=verbose,
+                **kwargs
+            )
+        else:
+            return self.encode_plus(
+                text=text,
+                text_pair=text_pair,
+                add_special_tokens=add_special_tokens,
+                padding=padding,
+                truncation=truncation,
+                max_length=max_length,
+                stride=stride,
+                is_split_into_words=is_split_into_words,
+                pad_to_multiple_of=pad_to_multiple_of,
+                return_tensors=return_tensors,
+                return_token_type_ids=return_token_type_ids,
+                return_attention_mask=return_attention_mask,
+                return_overflowing_tokens=return_overflowing_tokens,
+                return_special_tokens_mask=return_special_tokens_mask,
+                return_offsets_mapping=return_offsets_mapping,
+                return_length=return_length,
+                verbose=verbose,
+                **kwargs
+            )
+
+    def encode_plus(
+        self,
+        text: str,
+        text_pair: Optional[str] = None,
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        is_split_into_words: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, Any]] = None,
+        return_token_type_ids: Optional[bool] = True,
+        return_attention_mask: Optional[bool] = True,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        **kwargs
+    ) -> BatchEncoding:
+        if max_length is None: max_length = self.model_max_length
+        ids_a = list(self._cached_encode_str(text.strip()))
+        ids_b = list(self._cached_encode_str(text_pair.strip())) if text_pair else None
+
+        input_ids, token_type_ids = [], []
+        if add_special_tokens:
+            input_ids.append(self.bos_token_id); token_type_ids.append(0)
+            input_ids.extend(ids_a); token_type_ids.extend([0] * len(ids_a))
+            input_ids.append(self.eos_token_id); token_type_ids.append(0)
+            if ids_b is not None:
+                input_ids.extend(ids_b); token_type_ids.extend([1] * len(ids_b))
+                input_ids.append(self.eos_token_id); token_type_ids.append(1)
+        else:
+            input_ids = ids_a.copy(); token_type_ids = [0] * len(input_ids)
+            if ids_b is not None:
+                input_ids.extend(ids_b); token_type_ids.extend([1] * len(ids_b))
+
+        if truncation and len(input_ids) > max_length:
+            input_ids, token_type_ids = input_ids[:max_length], token_type_ids[:max_length]
+
+        encoded_dict = {"input_ids": input_ids}
+        if return_attention_mask:
+            if padding == True or padding == "max_length":
+                pad_len = max_length - len(input_ids)
+                if pad_len > 0:
+                    if self.padding_side == "right":
+                        input_ids.extend([self.pad_token_id] * pad_len)
+                        token_type_ids.extend([0] * pad_len)
+                    else:
+                        input_ids = [self.pad_token_id] * pad_len + input_ids
+                        token_type_ids = [0] * pad_len + token_type_ids
+            attention_mask = [0 if tid == self.pad_token_id else 1 for tid in input_ids]
+            encoded_dict["attention_mask"] = attention_mask
+        if return_token_type_ids: encoded_dict["token_type_ids"] = token_type_ids
+        if return_special_tokens_mask:
+            encoded_dict["special_tokens_mask"] = [
+                1 if tid in {self.bos_token_id, self.eos_token_id, self.pad_token_id, self.mask_token_id} else 0
+                for tid in input_ids
+            ]
+        if return_length:
+            encoded_dict["length"] = len([tid for tid in input_ids if tid != self.pad_token_id])
+
+        if return_tensors in ["pt", "torch"]:
+            out = {}
+            for k, v in encoded_dict.items():
+                if isinstance(v, list):
+                    tensor = torch.tensor(
+                        [self.unk_token_id if x is None else int(x) for x in v], dtype=torch.long
+                    ).unsqueeze(0)
+                    out[k] = tensor
+                else:
+                    out[k] = v
+            return BatchEncoding(out, tensor_type=return_tensors)
+        return BatchEncoding(encoded_dict, tensor_type=None)
+
+    def batch_encode_plus(
+        self,
+        batch_text_or_text_pairs: List[Union[str, Tuple[str, str]]],
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        is_split_into_words: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, Any]] = None,
+        return_token_type_ids: Optional[bool] = True,
+        return_attention_mask: Optional[bool] = True,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        **kwargs
+    ) -> BatchEncoding:
+        if padding is True: padding = "longest"
+        if padding == "max_length" and max_length is None: max_length = self.model_max_length
+
+        all_input_ids, all_token_type_ids, all_attention_masks = [], [], []
+        all_special_masks, all_lengths = [], []
+        for item in batch_text_or_text_pairs:
+            t, tp = item if isinstance(item, tuple) else (item, None)
+            enc = self.encode_plus(
+                text=t, text_pair=tp, add_special_tokens=add_special_tokens,
+                padding=False, truncation=truncation, max_length=max_length,
+                return_tensors=None, return_token_type_ids=return_token_type_ids,
+                return_attention_mask=return_attention_mask,
+                return_special_tokens_mask=return_special_tokens_mask,
+                return_length=return_length, **kwargs
+            )
+            ids, tt, am = enc["input_ids"], enc.get("token_type_ids", [0]*len(enc["input_ids"])), enc.get("attention_mask",[1]*len(enc["input_ids"]))
+            sm, ln = enc.get("special_tokens_mask",[0]*len(ids)), enc.get("length", len([x for x in ids if x != self.pad_token_id]))
+            all_input_ids.append(ids); all_token_type_ids.append(tt); all_attention_masks.append(am)
+            all_special_masks.append(sm); all_lengths.append(ln)
+
+        pad_to = max(len(x) for x in all_input_ids) if padding == "longest" else (max_length if padding == "max_length" else None)
+        batched = {
+            "input_ids": all_input_ids,
+            "token_type_ids": all_token_type_ids if return_token_type_ids else None,
+            "attention_mask": all_attention_masks if return_attention_mask else None,
+            "special_tokens_mask": all_special_masks if return_special_tokens_mask else None,
+            "length": all_lengths if return_length else None,
+        }
+        if pad_to is not None:
+            for key in ["input_ids","token_type_ids","attention_mask","special_tokens_mask"]:
+                if batched.get(key) is None: continue
+                padded = []
+                for seq in batched[key]:
+                    pad_len = pad_to - len(seq)
+                    pad_val = self.pad_token_id if key=="input_ids" else 0
+                    if pad_len > 0:
+                        seq = seq+[pad_val]*pad_len if self.padding_side=="right" else [pad_val]*pad_len+seq
+                    padded.append(seq)
+                batched[key] = padded
+
+        if return_tensors in ["pt", "torch"]:
+            def to_tensor(lst, pad_val=0):
+                return torch.tensor([[self.unk_token_id if x is None else int(x) for x in row] for row in lst], dtype=torch.long)
+            out = {}
+            if batched.get("input_ids") is not None: out["input_ids"] = to_tensor(batched["input_ids"], self.pad_token_id)
+            if batched.get("attention_mask") is not None: out["attention_mask"] = to_tensor(batched["attention_mask"],0)
+            if batched.get("token_type_ids") is not None: out["token_type_ids"] = to_tensor(batched["token_type_ids"],0)
+            if batched.get("special_tokens_mask") is not None: out["special_tokens_mask"] = to_tensor(batched["special_tokens_mask"],0)
+            if return_length and batched.get("length") is not None: out["length"] = torch.tensor([int(x) for x in batched["length"]], dtype=torch.long)
+            return BatchEncoding(out, tensor_type=return_tensors)
+        return BatchEncoding({k:v for k,v in batched.items() if v is not None}, tensor_type=None)
+
+    # ------------------------------
+    # Decoding
+    # ------------------------------
+    def decode(self, token_ids, skip_special_tokens=False, **kwargs):
+        if isinstance(token_ids, torch.Tensor): token_ids = token_ids.tolist()
+        special_ids = {self.bos_token_id,self.eos_token_id,self.pad_token_id,self.mask_token_id} if skip_special_tokens else set()
+        tokens = [self.id_to_token.get(tid,self.unk_token) for tid in token_ids if tid not in special_ids]
+        return "".join(tokens)
+
+    def batch_decode(self, sequences, skip_special_tokens=False, **kwargs):
+        if isinstance(sequences, torch.Tensor): sequences = sequences.tolist()
+        return [self.decode(seq, skip_special_tokens=skip_special_tokens, **kwargs) for seq in sequences]
+
+    def decode_with_trace(self, token_ids: List[int]):
+        print(f"\n🔍 Decoding {len(token_ids)} tokens:")
+        for i, tid in enumerate(token_ids):
+            token = self.id_to_token.get(tid, self.unk_token)
+            tid_str = "None" if tid is None else f"{tid:5d}"
+            print(f"  [{i:03d}] ID={tid_str} → '{token}'")
+
+    # ------------------------------
+    # Save / Load
+    # ------------------------------
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory): os.makedirs(save_directory)
+        vocab_file = os.path.join(save_directory,(filename_prefix+"-" if filename_prefix else "")+"vocab.json")
+        with open(vocab_file,"w",encoding="utf-8") as f: json.dump(self.token_to_id,f,ensure_ascii=False,indent=2)
+        return (vocab_file,)
+
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], filename_prefix: Optional[str]=None, **kwargs):
+        if not os.path.exists(save_directory): os.makedirs(save_directory)
+        self.save_vocabulary(save_directory, filename_prefix)
+        config_file = os.path.join(save_directory,"tokenizer_config.json")
+        with open(config_file,"w",encoding="utf-8") as f:
+            json.dump({
+                "tokenizer_class": self.__class__.__name__,
+                "model_max_length": self.model_max_length,
+                "padding_side": self.padding_side,
+                "truncation_side": self.truncation_side,
+                "special_tokens": {
+                    "bos_token": self.bos_token,
+                    "eos_token": self.eos_token,
+                    "pad_token": self.pad_token,
+                    "unk_token": self.unk_token,
+                    "mask_token": self.mask_token,
+                }
+            },f,ensure_ascii=False,indent=2)
+        return (save_directory,)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        if os.path.isdir(pretrained_model_name_or_path):
+            vocab_file = os.path.join(pretrained_model_name_or_path,"vocab.json")
+            config_file = os.path.join(pretrained_model_name_or_path,"tokenizer_config.json")
+            config = {}
+            if os.path.exists(config_file):
+                with open(config_file,"r",encoding="utf-8") as f: config=json.load(f)
+            return cls(vocab_file=vocab_file, **{**config,**kwargs})
+        else:
+            raise NotImplementedError("Loading from Hub not implemented yet")
+
+
+# ------------------------------
+# SELFIES variant
+# ------------------------------
+class FastChemTokenizerSelfies(FastChemTokenizer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)  # ✅ ensures BOS/EOS etc. are set
+
+    """SELFIES variant that handles whitespace-separated tokens."""
+
+    def _encode_core(self, text: str) -> List[int]:
+        result_ids, i, n = [], 0, len(text)
+        while i < n:
+            if text[i].isspace(): i += 1; continue
+            node, j = self.trie_root, i
+            last_match_id, last_match_end = None, i
+            while j < n and text[j] in node.children:
+                node = node.children[text[j]]; j += 1
+                if node.token_id is not None:
+                    last_match_id, last_match_end = node.token_id, j
+            if last_match_id is not None:
+                result_ids.append(last_match_id); i = last_match_end
+            else:
+                result_ids.append(self.token_to_id.get(text[i], self.unk_token_id)); i += 1
+        return result_ids
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str: return " ".join(tokens)
+    def decode(self, token_ids, skip_special_tokens=False, **kwargs):
+        if isinstance(token_ids, torch.Tensor): token_ids = token_ids.tolist()
+        special_ids = {self.bos_token_id,self.eos_token_id,self.pad_token_id,self.mask_token_id} if skip_special_tokens else set()
+        tokens = [self.id_to_token.get(tid,self.unk_token) for tid in token_ids if tid not in special_ids]
+        return " ".join(tokens)

ChemQ3MTP/__init__.py

@@ -0,0 +1,14 @@
+# __init__.py
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from .configuration_chemq3mtp import ChemQ3MTPConfig
+from .modeling_chemq3mtp import ChemQ3MTPForCausalLM
+from .FastChemTokenizerHF import FastChemTokenizerSelfies  
+
+# Register the model
+AutoConfig.register("chemq3_mtp", ChemQ3MTPConfig)
+AutoModelForCausalLM.register(ChemQ3MTPConfig, ChemQ3MTPForCausalLM)
+
+# Register the tokenizer
+AutoTokenizer.register(ChemQ3MTPConfig, FastChemTokenizerSelfies)
+
+__all__ = ["ChemQ3MTPConfig", "ChemQ3MTPForCausalLM", "FastChemTokenizerSelfies"]

ChemQ3MTP/configuration_chemq3mtp.py

@@ -0,0 +1,27 @@
+# configuration_chemq3mtp.py
+from transformers import Qwen2Config
+
+class ChemQ3MTPConfig(Qwen2Config):
+    """
+    Configuration class for ChemQ3MTP model.
+    """
+    model_type = "chemq3_mtp"
+    
+    def __init__(
+        self,
+        num_future_tokens: int = 3,
+        horizon_weights = None,
+        use_mtp_training: bool = True,
+        entropy_controller_config = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.num_future_tokens = num_future_tokens
+        self.horizon_weights = horizon_weights or [0.9 ** i for i in range(num_future_tokens)]
+        self.use_mtp_training = use_mtp_training
+        self.entropy_controller_config = entropy_controller_config or {
+            "min_entropy": 0.5,
+            "max_entropy": 3.0,
+            "target_entropy": 1.5,
+            "adaptation_rate": 0.01
+        }

ChemQ3MTP/misc_utils.py

@@ -0,0 +1,98 @@
+# src/misc_utils.py
+import json
+import torch
+import shutil
+import os
+import datetime
+from transformers import TrainerCallback
+
+def load_config(config_path: str = "config.json") -> dict:
+    with open(config_path, "r") as f:
+        return json.load(f)
+
+def get_training_config(config: dict) -> dict:
+    return config["training"]
+
+def get_model_config(config: dict) -> dict:
+    return {k: v for k, v in config.items() 
+             if k not in ["training", "generation", "model_type", "architectures"]}
+
+def get_generation_config(config: dict) -> dict:
+    return config.get("generation", {})
+
+def clear_cache():
+    print("Clearing PyTorch and CUDA caches...")
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        print("CUDA cache cleared")
+    torch.backends.cudnn.benchmark = True
+    print("PyTorch cache cleared")
+
+def clear_datasets_cache():
+    from datasets import get_cache_directory
+    try:
+        cache_dir = get_cache_directory()
+        print(f"Clearing datasets cache at: {cache_dir}")
+        if os.path.exists(cache_dir):
+            shutil.rmtree(cache_dir)
+            print("Datasets cache cleared")
+    except:
+        print("Could not clear datasets cache (may not exist)")
+
+class LossLoggerCallback(TrainerCallback):
+    def __init__(self, log_file="training_losses.txt", with_timestamp=False):
+        self.log_file = log_file
+        self.with_timestamp = with_timestamp
+        with open(self.log_file, "w") as f:
+            if self.with_timestamp:
+                f.write("time\tstep\tloss\teval_loss\n")
+            else:
+                f.write("step\tloss\teval_loss\n")
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs is None:
+            return
+        step = state.global_step
+        loss = logs.get("loss")
+        eval_loss = logs.get("eval_loss")
+
+        with open(self.log_file, "a") as f:
+            if self.with_timestamp:
+                ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                f.write(f"{ts}\t{step}\t{loss if loss is not None else ''}\t{eval_loss if eval_loss is not None else ''}\n")
+            else:
+                f.write(f"{step}\t{loss if loss is not None else ''}\t{eval_loss if eval_loss is not None else ''}\n")
+
+class CheckpointEvery10PercentCallback(TrainerCallback):
+    def __init__(self, save_dir, total_steps):
+        self.save_dir = save_dir
+        self.total_steps = total_steps
+        self.checkpoint_intervals = []
+        for i in range(1, 11):
+            checkpoint_step = int(total_steps * i * 0.1)
+            self.checkpoint_intervals.append(checkpoint_step)
+        self.saved_checkpoints = set()
+        print(f"Checkpoint intervals: {self.checkpoint_intervals}")
+
+    def on_step_end(self, args, state, control, **kwargs):
+        current_step = state.global_step
+        for checkpoint_step in self.checkpoint_intervals:
+            if current_step == checkpoint_step and checkpoint_step not in self.saved_checkpoints:
+                checkpoint_dir = f"{self.save_dir}/checkpoint_10percent_{current_step}"
+                print(f"Saving 10% progress checkpoint at step {current_step} to {checkpoint_dir}")
+                
+                model = kwargs.get('model')
+                tokenizer = kwargs.get('processing_class')
+                
+                if model is not None:
+                    model.save_pretrained(checkpoint_dir)
+                if tokenizer is not None:
+                    tokenizer.save_pretrained(checkpoint_dir)
+                
+                if hasattr(kwargs.get('trainer'), 'save_state'):
+                    kwargs['trainer'].save_state()
+                
+                self.saved_checkpoints.add(checkpoint_step)
+                print(f"Checkpoint saved at step {current_step} ({current_step/self.total_steps*100:.1f}% completion)")
+                break

ChemQ3MTP/modeling_chemq3mtp.py

@@ -0,0 +1,467 @@
+# ========================
+#  ChemQ3-MTP - HuggingFace Compatible Version
+#  MODEL COMPONENTS 
+#  by gbyuvd
+# ========================
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Categorical
+from typing import List, Union, Optional, Tuple, Dict, Any
+from transformers import Qwen2Config, Qwen2ForCausalLM, AutoTokenizer
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.utils import logging
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_utils import PreTrainedModel
+from rdkit import Chem
+from rdkit.Chem import Descriptors, Lipinski
+import selfies as sf
+from rdkit import RDLogger
+RDLogger.DisableLog('rdApp.*')
+import json
+import numpy as np
+from collections import Counter
+from rdkit.Chem import rdMolDescriptors
+
+logger = logging.get_logger(__name__)
+
+# ========================
+# CONFIGURATION CLASS
+# ========================
+
+class ChemQ3MTPConfig(Qwen2Config):
+    """
+    Configuration class for ChemQ3MTP model.
+    """
+    model_type = "chemq3_mtp"
+    
+    def __init__(
+        self,
+        num_future_tokens: int = 3,
+        horizon_weights: Optional[List[float]] = None,
+        use_mtp_training: bool = True,
+        entropy_controller_config: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.num_future_tokens = num_future_tokens
+        self.horizon_weights = horizon_weights or [0.9 ** i for i in range(num_future_tokens)]
+        self.use_mtp_training = use_mtp_training
+        self.entropy_controller_config = entropy_controller_config or {
+            "min_entropy": 0.5,
+            "max_entropy": 3.0,
+            "target_entropy": 1.5,
+            "adaptation_rate": 0.01
+        }
+
+# ========================
+# UTILITY FUNCTIONS (kept minimal for HF compatibility)
+# ========================
+
+def selfies_to_smiles(selfies_str: str) -> str | None:
+    """Convert SELFIES string to SMILES, handling tokenizer artifacts."""
+    try:
+        clean_selfies = selfies_str.replace(" ", "")
+        return sf.decoder(clean_selfies)
+    except Exception:
+        return None
+
+def is_valid_smiles(smiles: str) -> bool:
+    if not isinstance(smiles, str) or len(smiles.strip()) == 0:
+        return False
+    return Chem.MolFromSmiles(smiles.strip()) is not None
+
+# ========================
+# MODEL COMPONENTS
+# ========================
+
+class MTPHead(nn.Module):
+    """Multi-Token Prediction Head for predicting future tokens."""
+    
+    def __init__(self, hidden_size: int, vocab_size: int, num_future_tokens: int = 3):
+        super().__init__()
+        self.num_future_tokens = num_future_tokens
+        self.vocab_size = vocab_size
+        self.prediction_heads = nn.ModuleList([
+            nn.Linear(hidden_size, vocab_size, bias=False)
+            for _ in range(num_future_tokens)
+        ])
+        self.position_embeddings = nn.Embedding(num_future_tokens, hidden_size)
+        self.layer_norm = nn.LayerNorm(hidden_size)
+        
+    def forward(self, hidden_states: torch.Tensor) -> Dict[str, torch.Tensor]:
+        batch_size, seq_len, hidden_size = hidden_states.shape
+        outputs = {}
+        
+        for i in range(self.num_future_tokens):
+            pos_emb = self.position_embeddings(torch.tensor(i, device=hidden_states.device))
+            enhanced_hidden = self.layer_norm(hidden_states + pos_emb)
+            logits = self.prediction_heads[i](enhanced_hidden)
+            outputs[f'logits_t{i+1}'] = logits
+            
+        return outputs
+
+
+class HorizonLoss(nn.Module):
+    """Loss function for multi-horizon prediction."""
+    
+    def __init__(self, num_future_tokens: int = 3, horizon_weights: Optional[List[float]] = None):
+        super().__init__()
+        self.num_future_tokens = num_future_tokens
+        if horizon_weights is None:
+            self.horizon_weights = [0.9 ** i for i in range(num_future_tokens)]
+        else:
+            self.horizon_weights = horizon_weights
+        self.log_weights = nn.Parameter(torch.log(torch.tensor(self.horizon_weights)))
+        
+    def forward(
+        self, 
+        mtp_outputs: Dict[str, torch.Tensor], 
+        input_ids: torch.Tensor, 
+        attention_mask: Optional[torch.Tensor] = None
+    ) -> Dict[str, torch.Tensor]:
+        
+        batch_size, seq_len = input_ids.shape
+        device = input_ids.device
+        weights = F.softmax(self.log_weights, dim=0)
+        total_loss = 0.0
+        horizon_losses = {}
+        
+        for i in range(self.num_future_tokens):
+            logits_key = f'logits_t{i+1}'
+            if logits_key not in mtp_outputs:
+                continue
+                
+            logits = mtp_outputs[logits_key]
+            shift = i + 1
+            if seq_len <= shift:
+                continue
+                
+            shifted_logits = logits[:, :-shift, :].contiguous()
+            shifted_targets = input_ids[:, shift:].contiguous()
+            
+            if attention_mask is not None:
+                shifted_mask = attention_mask[:, shift:].contiguous()
+                mask_expanded = shifted_mask.view(-1)
+                valid_indices = mask_expanded == 1
+                if valid_indices.sum() == 0:
+                    continue
+                flat_logits = shifted_logits.view(-1, logits.size(-1))[valid_indices]
+                flat_targets = shifted_targets.view(-1)[valid_indices]
+            else:
+                flat_logits = shifted_logits.view(-1, logits.size(-1))
+                flat_targets = shifted_targets.view(-1)
+                
+            horizon_loss = F.cross_entropy(flat_logits, flat_targets, reduction='mean')
+            horizon_losses[f'horizon_loss_t{i+1}'] = horizon_loss
+            total_loss += weights[i] * horizon_loss
+            
+        return {'loss': total_loss, 'horizon_weights': weights, **horizon_losses}
+
+
+class EnhancedEntropyController:
+    """Enhanced entropy controller for adaptive training."""
+    
+    def __init__(self, min_entropy: float = 0.5, max_entropy: float = 3.0,
+                 target_entropy: float = 1.5, adaptation_rate: float = 0.01):
+        self.min_entropy = min_entropy
+        self.max_entropy = max_entropy
+        self.target_entropy = target_entropy
+        self.adaptation_rate = adaptation_rate
+        self.entropy_history = []
+        self.entropy_weight = 0.01
+        
+    def update_entropy_weight(self, current_entropy: float) -> float:
+        """Dynamically adjust entropy weight based on current entropy levels."""
+        self.entropy_history.append(current_entropy)
+        
+        if len(self.entropy_history) > 100:
+            self.entropy_history = self.entropy_history[-100:]
+            
+        if len(self.entropy_history) >= 10:
+            avg_entropy = np.mean(self.entropy_history[-10:])
+            
+            if avg_entropy < self.target_entropy * 0.8:
+                self.entropy_weight = min(0.05, self.entropy_weight * 1.1)
+            elif avg_entropy > self.target_entropy * 1.2:
+                self.entropy_weight = max(0.001, self.entropy_weight * 0.95)
+                
+        return self.entropy_weight
+
+# ========================
+# MAIN MODEL CLASS
+# ========================
+
+class ChemQ3MTPForCausalLM(Qwen2ForCausalLM):
+    """
+    ChemQ3MTP model for causal language modeling with multi-token prediction.
+    
+    This model extends Qwen2ForCausalLM with additional capabilities for
+    multi-token prediction and chemistry-specific training.
+    """
+    
+    config_class = ChemQ3MTPConfig
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+    
+    def __init__(self, config: ChemQ3MTPConfig):
+        super().__init__(config)
+        
+        # Initialize MTP components
+        self.mtp_head = MTPHead(
+            config.hidden_size, 
+            config.vocab_size, 
+            config.num_future_tokens
+        )
+        self.horizon_loss = HorizonLoss(
+            num_future_tokens=config.num_future_tokens,
+            horizon_weights=config.horizon_weights
+        )
+        
+        # Training configuration
+        self.use_mtp_training = config.use_mtp_training
+        
+        # Initialize entropy controller
+        self.entropy_controller = EnhancedEntropyController(
+            **config.entropy_controller_config
+        )
+        
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        """
+        Forward pass of the ChemQ3MTP model.
+        """
+        
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Default attention mask if not provided
+        if attention_mask is None and input_ids is not None:
+            # Handle case where pad_token_id is None
+            if hasattr(self.config, 'pad_token_id') and self.config.pad_token_id is not None:
+                attention_mask = (input_ids != self.config.pad_token_id).long()
+            else:
+                # Default to all 1s if no pad_token_id is defined
+                attention_mask = torch.ones_like(input_ids, dtype=torch.long)
+
+        # Call parent forward with required hidden states
+        outputs = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=None,  # Handle labels manually
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=True,  # Always need hidden states for MTP
+            return_dict=True,
+            cache_position=cache_position,
+        )
+
+        # Rest of your forward method...
+        hidden_states = outputs.hidden_states[-1]
+        lm_logits = outputs.logits
+        loss = None
+
+        # Compute loss if labels are provided
+        if labels is not None:
+            if self.training and self.use_mtp_training:
+                # Multi-token prediction training
+                mtp_outputs = self.mtp_head(hidden_states)
+                horizon_loss_dict = self.horizon_loss(mtp_outputs, input_ids, attention_mask)
+
+                # Standard causal LM loss
+                shift_logits = lm_logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+
+                if attention_mask is not None:
+                    shift_mask = attention_mask[..., 1:].contiguous()
+                    loss_mask = shift_mask.view(-1) == 1
+                    if loss_mask.sum() == 0:
+                        causal_lm_loss = torch.tensor(0.0, device=lm_logits.device)
+                    else:
+                        flat_logits = shift_logits.view(-1, shift_logits.size(-1))[loss_mask]
+                        flat_labels = shift_labels.view(-1)[loss_mask]
+                        causal_lm_loss = F.cross_entropy(flat_logits, flat_labels, reduction='mean')
+                else:
+                    flat_logits = shift_logits.view(-1, shift_logits.size(-1))
+                    flat_labels = shift_labels.view(-1)
+                    causal_lm_loss = F.cross_entropy(flat_logits, flat_labels, reduction='mean')
+
+                # Combine losses
+                loss = 0.7 * horizon_loss_dict['loss'] + 0.3 * causal_lm_loss
+
+            else:
+                # Standard causal LM training
+                shift_logits = lm_logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                loss = F.cross_entropy(
+                    shift_logits.view(-1, shift_logits.size(-1)),
+                    shift_labels.view(-1),
+                    ignore_index=-100
+                )
+
+        if not return_dict:
+            output = (lm_logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def set_mtp_training(self, use_mtp: bool):
+        """Enable or disable multi-token prediction training."""
+        self.use_mtp_training = use_mtp
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        **kwargs
+    ):
+        """
+        Prepare inputs for generation. This method is required for compatibility
+        with HuggingFace's generation utilities.
+        """
+        # This delegates to the parent class implementation
+        return super().prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            **kwargs
+        )
+
+    def generate_with_logprobs(
+        self,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int = 50,
+        temperature: float = 1.0,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        do_sample: bool = True,
+        return_probs: bool = True,
+        tokenizer=None,
+    ) -> Tuple[List[str], torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Generate sequences with log probabilities for RL training.
+        
+        FIXED VERSION: Corrects log probability calculation to avoid numerical issues.
+        Changes:
+        1. Use log_softmax instead of log(softmax) to avoid log(0) issues
+        2. Correct the gather operation for non-sampling case
+        3. Handle the case where filtered logits become -inf properly
+        """
+        self.eval()
+        device = input_ids.device
+
+        # Normalize input shapes
+        if input_ids.dim() == 1:
+            input_ids = input_ids.unsqueeze(0)
+        if input_ids.dim() == 3 and input_ids.size(1) == 1:
+            input_ids = input_ids.squeeze(1)
+        assert input_ids.dim() == 2, f"input_ids must be 2-D, got {input_ids.shape}"
+
+        batch_size, seq_len = input_ids.shape
+        current_input = input_ids
+
+        generated_tokens, generated_logprobs, generated_probs = [], [], []
+
+        with torch.no_grad():
+            for _ in range(max_new_tokens):
+                outputs = self(current_input, use_cache=False)
+                logits = outputs.logits[:, -1, :] / temperature
+
+                # Apply top-k filtering
+                if top_k is not None:
+                    values, indices = torch.topk(logits, k=top_k)
+                    logits = torch.full_like(logits, float("-inf"))
+                    logits.scatter_(1, indices, values)
+
+                # Apply top-p filtering
+                if top_p is not None and top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                    cumprobs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                    mask = cumprobs > top_p
+                    mask[..., 1:] = mask[..., :-1].clone()
+                    mask[..., 0] = False
+                    logits[mask.scatter(1, sorted_indices, mask)] = float("-inf")
+
+                # FIX: Calculate log probabilities using log_softmax for numerical stability
+                log_probs = F.log_softmax(logits, dim=-1)
+                probs = F.softmax(logits, dim=-1)
+
+                if do_sample:
+                    dist = Categorical(probs)
+                    next_token = dist.sample()
+                    # FIX: Get log prob directly from log_probs tensor
+                    log_p = torch.gather(log_probs, 1, next_token.unsqueeze(1)).squeeze(1)
+                else:
+                    next_token = torch.argmax(probs, dim=-1)
+                    # FIX: Use log_probs instead of log(probs) to avoid numerical issues
+                    log_p = torch.gather(log_probs, 1, next_token.unsqueeze(1)).squeeze(1)
+
+                generated_tokens.append(next_token.unsqueeze(1))
+                generated_logprobs.append(log_p.unsqueeze(1))
+                if return_probs:
+                    generated_probs.append(probs.unsqueeze(1))
+
+                current_input = torch.cat([current_input, next_token.unsqueeze(1)], dim=1)
+
+        generated_tokens = torch.cat(generated_tokens, dim=1)
+        generated_logprobs = torch.cat(generated_logprobs, dim=1)
+        generated_probs = torch.cat(generated_probs, dim=1) if return_probs else None
+
+        # Decode generated tokens
+        if tokenizer is None:
+            tokenizer = getattr(self, "tokenizer", None)
+        if tokenizer is None:
+            raise ValueError("Tokenizer must be provided to decode generated tokens.")
+
+        decoded_list = [
+            tokenizer.decode(tok_ids, skip_special_tokens=True)
+            for tok_ids in generated_tokens
+        ]
+        
+        return decoded_list, generated_logprobs, generated_tokens, generated_probs
+
+# ========================
+# REGISTRATION
+# ========================
+
+# Register the configuration and model classes
+from transformers import AutoConfig, AutoModelForCausalLM
+
+# Register the configuration and model classes
+AutoConfig.register("chemq3_mtp", ChemQ3MTPConfig)
+AutoModelForCausalLM.register(ChemQ3MTPConfig, ChemQ3MTPForCausalLM)

ChemQ3MTP/rl_utils.py

@@ -0,0 +1,1070 @@
+# ========================
+#  RL_UTILS.PY
+#  v3
+#  Chemistry RL Training Utilities for ChemQ3-MTP
+#  by gbyuvd
+#  Patched: reward normalization, KL/entropy reset per phase,
+#           entropy target annealing, and symmetric curriculum
+#           and now with Durrant's Lab's filtering included
+# ========================
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Categorical
+from typing import List, Union, Optional, Tuple, Dict, Any
+import numpy as np
+from collections import Counter, deque
+
+# Chemistry imports
+from rdkit import Chem
+from rdkit.Chem import Descriptors, Lipinski, rdMolDescriptors
+import selfies as sf
+from rdkit import RDLogger
+RDLogger.DisableLog('rdApp.*')
+
+# Optional: HuggingFace for SA classifier
+try:
+    from transformers import pipeline, AutoTokenizer
+    HF_AVAILABLE = True
+except ImportError:
+    HF_AVAILABLE = False
+    print("Warning: transformers not available, SA classifier will not work")
+
+# ========================
+# CHEMISTRY UTILITIES
+# ========================
+
+def selfies_to_smiles(selfies_str: str) -> str | None:
+    """Convert SELFIES string to SMILES, handling tokenizer artifacts."""
+    try:
+        clean_selfies = selfies_str.replace(" ", "")
+        return sf.decoder(clean_selfies)
+    except Exception:
+        return None
+
+def is_valid_smiles(smiles: str) -> bool:
+    """
+    Check if a SMILES string represents a valid molecule.
+    FIXED: Now properly checks for heavy atoms (non-hydrogens) >= 3
+    and rejects disconnected/separated molecules
+    """
+    if not isinstance(smiles, str) or len(smiles.strip()) == 0:
+        return False
+    
+    smiles = smiles.strip()
+    
+    # FAST CHECK: Reject separated molecules (contains dots)
+    if '.' in smiles:
+        return False  # Disconnected components indicated by dots
+    
+    try:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return False
+        
+        # CRITICAL FIX: Check heavy atoms (non-hydrogens), not total atoms
+        heavy_atoms = mol.GetNumHeavyAtoms()  # This excludes hydrogens
+        if heavy_atoms < 3:
+            return False
+            
+        return True
+    except Exception:
+        return False
+
+def passes_durrant_lab_filter(smiles: str) -> bool:
+    """
+    Apply Durrant's lab filter to remove improbable substructures.
+    FIXED: More robust error handling, pattern checking, and disconnected molecule rejection.
+    Returns True if molecule passes the filter (is acceptable), False otherwise.
+    """
+    if not smiles or not isinstance(smiles, str) or len(smiles.strip()) == 0:
+        return False
+    
+    try:
+        mol = Chem.MolFromSmiles(smiles.strip())
+        if mol is None:
+            return False
+        
+        # Check heavy atoms again (belt and suspenders approach)
+        if mol.GetNumHeavyAtoms() < 3:
+            return False
+        
+        # REJECT SEPARATED/DISCONNECTED MOLECULES (double check here too)
+        fragments = Chem.rdmolops.GetMolFrags(mol, asMols=False)
+        if len(fragments) > 1:
+            return False  # Reject molecules with multiple disconnected parts
+        
+        # Define SMARTS patterns for problematic substructures
+        problematic_patterns = [
+            "C=[N-]",                    # Carbon double bonded to negative nitrogen
+            "[N-]C=[N+]",               # Nitrogen anion bonded to nitrogen cation
+            "[nH+]c[n-]",               # Aromatic nitrogen cation adjacent to nitrogen anion
+            "[#7+]~[#7+]",              # Positive nitrogen connected to positive nitrogen
+            "[#7-]~[#7-]",              # Negative nitrogen connected to negative nitrogen
+            "[!#7]~[#7+]~[#7-]~[!#7]",  # Bridge: non-nitrogen - pos nitrogen - neg nitrogen - non-nitrogen
+            "[#5]",                     # Boron atoms
+            "O=[PH](=O)([#8])([#8])",   # Phosphoryl with hydroxyls
+            "N=c1cc[#7]c[#7]1",         # Nitrogen in aromatic ring with another nitrogen
+            "[$([NX2H1]),$([NX3H2])]=C[$([OH]),$([O-])]", # N=CH-OH or N=CH-O-
+        ]
+        
+        # Check for metals (excluding common biologically relevant ions)
+        metal_exclusions = {11, 12, 19, 20}  # Na, Mg, K, Ca
+        for atom in mol.GetAtoms():
+            atomic_num = atom.GetAtomicNum()
+            # More precise metal detection
+            if atomic_num > 20 and atomic_num not in metal_exclusions:
+                return False
+        
+        # Check for each problematic pattern
+        for pattern in problematic_patterns:
+            try:
+                patt_mol = Chem.MolFromSmarts(pattern)
+                if patt_mol is not None:
+                    matches = mol.GetSubstructMatches(patt_mol)
+                    if matches:
+                        return False  # Found problematic substructure
+            except Exception:
+                # If SMARTS parsing fails, continue to next pattern
+                continue
+        
+        return True  # Passed all checks
+        
+    except Exception:
+        return False
+# ========================
+# SA CLASSIFIER
+# ========================
+
+# Global classifier instance for lazy loading
+_sa_classifier = None
+
+def get_sa_classifier():
+    """Get or initialize the synthetic accessibility classifier."""
+    global _sa_classifier
+    if not HF_AVAILABLE:
+        raise ImportError("transformers package required for SA classifier")
+    
+    if _sa_classifier is None:
+        try:
+            sa_tokenizer = AutoTokenizer.from_pretrained("gbyuvd/synthaccess-chemselfies")
+            _sa_classifier = pipeline(
+                "text-classification",
+                model="gbyuvd/synthaccess-chemselfies",
+                tokenizer=sa_tokenizer
+            )
+        except Exception as e:
+            print(f"Warning: Could not load SA classifier: {e}")
+            return None
+    return _sa_classifier
+
+def compute_sa_reward(selfies_str: str) -> float:
+    """Reward molecules with easy synthetic accessibility (SA)."""
+    try:
+        classifier = get_sa_classifier()
+        if classifier is None:
+            return 0.0
+        
+        result = classifier(selfies_str, truncation=True, max_length=128)[0]
+        if result["label"].lower() == "easy":
+            return result["score"]
+        else:
+            return -result["score"]  # penalize "Hard"
+    except Exception:
+        return 0.0
+    
+
+# ========================
+# MOLECULAR REWARD COMPONENTS
+# ========================
+
+def compute_biological_diversity_score(mol) -> float:
+    """Reward molecules with diverse CHONP atoms, normalized to [0,1]."""
+    if mol is None:
+        return 0.0
+    try:
+        atoms = [atom.GetSymbol() for atom in mol.GetAtoms()]
+        atom_counts = Counter(atoms)
+        bio_elements = {"C", "H", "O", "N", "P"}
+        present_bio_elements = set(atoms) & bio_elements
+
+        if len(present_bio_elements) < 2:
+            return 0.0
+
+        base_score = 0.3
+        diversity_bonus = (len(present_bio_elements) - 2) / 3 * 0.4
+
+        total_bio_atoms = sum(atom_counts.get(e, 0) for e in present_bio_elements)
+        if total_bio_atoms > 0:
+            bio_probs = [atom_counts.get(e, 0) / total_bio_atoms for e in present_bio_elements]
+            if len(bio_probs) > 1:
+                entropy = -sum(p * np.log2(p) for p in bio_probs if p > 0)
+                max_entropy = np.log2(len(bio_probs))
+                entropy_bonus = (entropy / max_entropy) * 0.3
+            else:
+                entropy_bonus = 0.0
+        else:
+            entropy_bonus = 0.0
+
+        return min(1.0, base_score + diversity_bonus + entropy_bonus)
+    except Exception:
+        return 0.0
+
+def compute_charge_neutrality_score(mol) -> float:
+    """Reward if molecule is globally neutral (formal charge = 0)."""
+    if mol is None:
+        return 0.0
+    try:
+        return 1.0 if Chem.rdmolops.GetFormalCharge(mol) == 0 else 0.0
+    except Exception:
+        return 0.0
+
+def compute_local_charge_penalty(mol) -> float:
+    """
+    Penalize carbocations/anions.
+    Returns 1.0 if no charged atoms, decreases with fraction charged.
+    """
+    if mol is None:
+        return 0.0
+    try:
+        charges = [atom.GetFormalCharge() for atom in mol.GetAtoms()]
+        if not charges:
+            return 1.0
+        charged_atoms = sum(1 for c in charges if c != 0)
+        total_atoms = len(charges)
+        return max(0.0, 1.0 - (charged_atoms / total_atoms))
+    except Exception:
+        return 0.0
+
+def compute_enhanced_lipinski_reward(mol) -> float:
+    """Soft Lipinski scoring with partial credit."""
+    if mol is None:
+        return 0.0
+    try:
+        mw = Descriptors.MolWt(mol)
+        logp = Descriptors.MolLogP(mol)
+        hbd = Lipinski.NumHDonors(mol)
+        hba = Lipinski.NumHAcceptors(mol)
+        scores = []
+
+        # Molecular Weight
+        if 250 <= mw <= 500:
+            scores.append(1.0)
+        elif 150 <= mw < 250:
+            scores.append(0.5)
+        elif 500 < mw <= 600:
+            scores.append(0.7)
+        else:
+            scores.append(0.0)
+
+        # LogP
+        if -1 <= logp <= 5:
+            scores.append(1.0)
+        elif -2 <= logp < -1 or 5 < logp <= 6:
+            scores.append(0.5)
+        else:
+            scores.append(0.0)
+
+        # Hydrogen bond donors
+        scores.append(1.0 if hbd <= 5 else max(0.0, 1.0 - 0.2 * (hbd - 5)))
+        
+        # Hydrogen bond acceptors
+        scores.append(1.0 if hba <= 10 else max(0.0, 1.0 - 0.1 * (hba - 10)))
+
+        return sum(scores) / len(scores)
+    except Exception:
+        return 0.0
+
+def compute_structural_complexity_reward(mol) -> float:
+    """Reward moderate complexity: 1–3 rings and some flexibility."""
+    if mol is None:
+        return 0.0
+    try:
+        ring_count = rdMolDescriptors.CalcNumRings(mol)
+        if 1 <= ring_count <= 3:
+            ring_score = 1.0
+        elif ring_count == 0:
+            ring_score = 0.3
+        elif ring_count <= 5:
+            ring_score = 0.7
+        else:
+            ring_score = 0.1
+
+        rot_bonds = Descriptors.NumRotatableBonds(mol)
+        if 2 <= rot_bonds <= 8:
+            flex_score = 1.0
+        elif rot_bonds <= 12:
+            flex_score = 0.7
+        elif rot_bonds in (0, 1):
+            flex_score = 0.5
+        else:
+            flex_score = 0.2
+
+        return (ring_score + flex_score) / 2
+    except Exception:
+        return 0.0
+
+def compute_lipinski_reward(mol) -> float:
+    """Simple Lipinski rule compliance scoring."""
+    if mol is None:
+        return 0.0
+    try:
+        mw = Descriptors.MolWt(mol)
+        logp = Descriptors.MolLogP(mol)
+        hbd = Lipinski.NumHDonors(mol)
+        hba = Lipinski.NumHAcceptors(mol)
+        
+        # We don't want too small fragments, so MW > 250
+        rules = [250 < mw <= 500, logp <= 5, hbd <= 5, hba <= 10]
+        return sum(rules) / 4.0
+    except Exception:
+        return 0.0
+
+# ========================
+# COMPREHENSIVE REWARD SYSTEM
+# ========================
+
+def compute_comprehensive_reward(selfies_str: str) -> Dict[str, float]:
+    """
+    Compute comprehensive reward for a SELFIES string.
+    
+    Args:
+        selfies_str: SELFIES representation of molecule
+        
+    Returns:
+        Dictionary containing individual reward components and total
+    """
+    smiles = selfies_to_smiles(selfies_str)
+    
+    # Check validity first
+    is_valid = (smiles is not None and 
+                is_valid_smiles(smiles) and 
+                passes_durrant_lab_filter(smiles))
+    
+    if is_valid:
+        mol = Chem.MolFromSmiles(smiles)
+    else:
+        mol = None
+
+    rewards = {
+        "validity": 1.0 if is_valid else 0.0,
+        "biological_diversity": compute_biological_diversity_score(mol),
+        "charge_neutrality": compute_charge_neutrality_score(mol),
+        "local_charge_penalty": compute_local_charge_penalty(mol),
+        "lipinski": compute_enhanced_lipinski_reward(mol),
+        "structural_complexity": compute_structural_complexity_reward(mol),
+    }
+
+    if not is_valid:
+        # If not valid, set all chemistry-based rewards to 0
+        for key in rewards:
+            if key != "validity":
+                rewards[key] = 0.0
+        rewards["total"] = 0.0
+    else:
+        # Weighted combination of rewards
+        weights = {
+            "validity": 1.0,
+            "biological_diversity": 2.0,
+            "charge_neutrality": 1.5,
+            "local_charge_penalty": 1.0,
+            "lipinski": 1.0,
+            "structural_complexity": 0.5,
+        }
+        weighted_sum = sum(rewards[k] * weights[k] for k in weights)
+        rewards["total"] = weighted_sum / sum(weights.values())
+
+    return rewards
+    
+def selfies_to_lipinski_reward(selfies_str: str) -> float:
+    """Convert SELFIES to SMILES, then compute Lipinski reward."""
+    smiles = selfies_to_smiles(selfies_str)
+    if smiles is None or not is_valid_smiles(smiles) or not passes_durrant_lab_filter(smiles):
+        return 0.0
+    mol = Chem.MolFromSmiles(smiles)
+    return compute_lipinski_reward(mol)
+
+# ========================
+# PARETO-STYLE DYNAMIC REWARD CONTROLLER
+# ========================
+
+class ParetoRewardController:
+    """
+    Dynamic reward mixing based on Pareto optimality principles.
+    Adapts reward weights based on current population performance.
+    """
+    
+    def __init__(
+        self,
+        objectives: List[str] = None,
+        history_size: int = 500,
+        adaptation_rate: float = 0.1,
+        min_weight: float = 0.05,
+        max_weight: float = 0.95,
+        pareto_pressure: float = 1.0,
+        exploration_phase_length: int = 100
+    ):
+        """
+        Args:
+            objectives: List of objective names to track
+            history_size: Size of rolling history for Pareto analysis
+            adaptation_rate: How quickly weights adapt (0-1)
+            min_weight: Minimum weight for any objective
+            max_weight: Maximum weight for any objective  
+            pareto_pressure: Higher = more aggressive toward Pareto front
+            exploration_phase_length: Steps of pure exploration before adaptation
+        """
+        self.objectives = objectives or ["total", "sa", "validity", "diversity"]
+        self.history_size = history_size
+        self.adaptation_rate = adaptation_rate
+        self.min_weight = min_weight
+        self.max_weight = max_weight
+        self.pareto_pressure = pareto_pressure
+        self.exploration_phase_length = exploration_phase_length
+        
+        # Initialize weights equally
+        n_objectives = len(self.objectives)
+        self.weights = {obj: 1.0/n_objectives for obj in self.objectives}
+        
+        # History tracking
+        self.objective_history = deque(maxlen=history_size)
+        self.pareto_history = deque(maxlen=100)  # Track Pareto front evolution
+        self.step_count = 0
+        
+        # Performance tracking
+        self.objective_trends = {obj: deque(maxlen=50) for obj in self.objectives}
+        self.stagnation_counters = {obj: 0 for obj in self.objectives}
+        
+    def update(self, batch_objectives: Dict[str, torch.Tensor]) -> Dict[str, float]:
+        """
+        Update weights based on current batch performance.
+        
+        Args:
+            batch_objectives: Dict of objective_name -> tensor of scores
+            
+        Returns:
+            Updated weights dictionary
+        """
+        self.step_count += 1
+        
+        # Convert to numpy for easier manipulation
+        batch_data = {}
+        for obj_name, tensor_vals in batch_objectives.items():
+            if obj_name in self.objectives:
+                batch_data[obj_name] = tensor_vals.detach().cpu().numpy()
+        
+        # Store in history
+        if len(batch_data) > 0:
+            batch_size = len(batch_data[next(iter(batch_data))])
+            for i in range(batch_size):
+                point = {obj: batch_data[obj][i] for obj in self.objectives if obj in batch_data}
+                if len(point) == len(self.objectives):  # Only store complete points
+                    self.objective_history.append(point)
+        
+        # Skip adaptation during exploration phase
+        if self.step_count <= self.exploration_phase_length:
+            return self.weights.copy()
+        
+        # Compute current Pareto front
+        current_front = self._compute_pareto_front()
+        if len(current_front) > 0:
+            self.pareto_history.append(len(current_front))
+        
+        # Adapt weights based on multiple criteria
+        self._adapt_weights_pareto_driven(batch_data)
+        self._adapt_weights_stagnation_driven(batch_data)
+        self._adapt_weights_diversity_driven()
+        
+        # Ensure constraints
+        self._normalize_weights()
+        
+        return self.weights.copy()
+    
+    def _compute_pareto_front(self) -> List[Dict[str, float]]:
+        """Compute current Pareto front from history."""
+        if len(self.objective_history) < 10:
+            return []
+        
+        points = list(self.objective_history)
+        pareto_front = []
+        
+        for i, point1 in enumerate(points):
+            is_dominated = False
+            for j, point2 in enumerate(points):
+                if i != j and self._dominates(point2, point1):
+                    is_dominated = True
+                    break
+            if not is_dominated:
+                pareto_front.append(point1)
+        
+        return pareto_front
+    
+    def _dominates(self, point1: Dict[str, float], point2: Dict[str, float]) -> bool:
+        """Check if point1 dominates point2 (higher is better for all objectives)."""
+        better_in_all = True
+        strictly_better_in_one = False
+        
+        for obj in self.objectives:
+            if obj in point1 and obj in point2:
+                if point1[obj] < point2[obj]:
+                    better_in_all = False
+                    break
+                elif point1[obj] > point2[obj]:
+                    strictly_better_in_one = True
+        
+        return better_in_all and strictly_better_in_one
+    
+    def _adapt_weights_pareto_driven(self, batch_data: Dict[str, np.ndarray]):
+        """Adapt weights based on distance to Pareto front."""
+        if len(self.objective_history) < 50:
+            return
+        
+        pareto_front = self._compute_pareto_front()
+        if len(pareto_front) == 0:
+            return
+        
+        # Compute average distance to Pareto front for each objective
+        obj_distances = {obj: [] for obj in self.objectives}
+        
+        for point in list(self.objective_history)[-100:]:  # Recent history
+            min_distance = float('inf')
+            closest_front_point = None
+            
+            for front_point in pareto_front:
+                distance = sum((point[obj] - front_point[obj])**2 
+                             for obj in self.objectives if obj in point and obj in front_point)
+                if distance < min_distance:
+                    min_distance = distance
+                    closest_front_point = front_point
+            
+            if closest_front_point:
+                for obj in self.objectives:
+                    if obj in point and obj in closest_front_point:
+                        obj_distances[obj].append(abs(point[obj] - closest_front_point[obj]))
+        
+        # Increase weight for objectives with larger gaps to Pareto front
+        for obj in self.objectives:
+            if obj_distances[obj]:
+                avg_distance = np.mean(obj_distances[obj])
+                # Higher distance = increase weight
+                weight_adjustment = avg_distance * self.adaptation_rate * self.pareto_pressure
+                self.weights[obj] = self.weights[obj] * (1 + weight_adjustment)
+    
+    def _adapt_weights_stagnation_driven(self, batch_data: Dict[str, np.ndarray]):
+        """Increase weights for stagnating objectives."""
+        for obj in self.objectives:
+            if obj in batch_data:
+                current_mean = np.mean(batch_data[obj])
+                self.objective_trends[obj].append(current_mean)
+                
+                if len(self.objective_trends[obj]) >= 20:
+                    recent_trend = np.array(list(self.objective_trends[obj])[-20:])
+                    # Check for stagnation (low variance)
+                    if np.std(recent_trend) < 0.01:  # Adjust threshold as needed
+                        self.stagnation_counters[obj] += 1
+                        # Boost weight for stagnating objectives
+                        boost = min(0.1, self.stagnation_counters[obj] * 0.02)
+                        self.weights[obj] += boost
+                    else:
+                        self.stagnation_counters[obj] = max(0, self.stagnation_counters[obj] - 1)
+    
+    def _adapt_weights_diversity_driven(self):
+        """Adapt weights based on Pareto front diversity."""
+        if len(self.pareto_history) < 10:
+            return
+        
+        recent_front_sizes = list(self.pareto_history)[-10:]
+        front_diversity = np.std(recent_front_sizes)
+        
+        # If diversity is low, boost exploration objectives
+        if front_diversity < 1.0:  # Adjust threshold
+            exploration_objectives = ["sa", "diversity"]  # Objectives that promote exploration
+            for obj in exploration_objectives:
+                if obj in self.weights:
+                    self.weights[obj] += 0.05 * self.adaptation_rate
+    
+    def _normalize_weights(self):
+        """Ensure weights are normalized and within bounds."""
+        # Apply bounds
+        for obj in self.weights:
+            self.weights[obj] = np.clip(self.weights[obj], self.min_weight, self.max_weight)
+        
+        # Normalize to sum to 1
+        total = sum(self.weights.values())
+        if total > 0:
+            for obj in self.weights:
+                self.weights[obj] /= total
+        else:
+            # Fallback to equal weights
+            n = len(self.weights)
+            for obj in self.weights:
+                self.weights[obj] = 1.0 / n
+
+    def get_mixed_reward(self, rewards_dict: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """
+        Compute mixed reward using current weights.
+        
+        Args:
+            rewards_dict: Dictionary of reward tensors
+            
+        Returns:
+            Mixed reward tensor
+        """
+        mixed_reward = None
+        
+        for obj_name, weight in self.weights.items():
+            if obj_name in rewards_dict:
+                weighted_reward = weight * rewards_dict[obj_name]
+                if mixed_reward is None:
+                    mixed_reward = weighted_reward
+                else:
+                    mixed_reward += weighted_reward
+        
+        return mixed_reward if mixed_reward is not None else torch.zeros_like(list(rewards_dict.values())[0])
+    
+    def get_status(self) -> Dict[str, any]:
+        """Get current status for logging."""
+        pareto_front = self._compute_pareto_front()
+        
+        return {
+            "weights": self.weights.copy(),
+            "step_count": self.step_count,
+            "pareto_front_size": len(pareto_front),
+            "stagnation_counters": self.stagnation_counters.copy(),
+            "history_size": len(self.objective_history),
+            "avg_pareto_size": np.mean(list(self.pareto_history)) if self.pareto_history else 0
+        }
+
+
+# ========================
+# RL TRAINING CONTROLLERS
+# ========================
+
+class AdaptiveKLController:
+    """
+    Adaptive KL controller with hard clipping and EMA smoothing.
+    Prevents runaway beta values and exploding KL penalties.
+    """
+
+    def __init__(
+        self,
+        init_kl_coef: float = 0.2,
+        target_kl: float = 6.0,
+        horizon: int = 10000,
+        max_kl_coef: float = 10.0,
+        max_inc_factor: float = 2.0,
+        ema_alpha: float = 0.9,
+        kl_penalty_cap: float = 10.0,
+    ):
+        self.value = init_kl_coef
+        self.target = target_kl
+        self.horizon = horizon
+        self.max_kl_coef = max_kl_coef
+        self.max_inc_factor = max_inc_factor
+        self.ema_alpha = ema_alpha
+        self.kl_penalty_cap = kl_penalty_cap
+
+        # Exponential moving average of KL
+        self.ema_kl = None
+
+    def update(self, current_kl: float, n_steps: int) -> None:
+        # update EMA
+        if self.ema_kl is None:
+            self.ema_kl = current_kl
+        else:
+            self.ema_kl = (
+                self.ema_alpha * self.ema_kl + (1 - self.ema_alpha) * current_kl
+            )
+
+        proportional_error = np.clip(
+            (self.ema_kl - self.target) / self.target, -1.0, 1.0
+        )
+        mult = 1.0 + proportional_error * (n_steps / self.horizon)
+
+        # cap growth
+        if mult > self.max_inc_factor:
+            mult = self.max_inc_factor
+
+        # update beta
+        new_val = self.value * mult
+        self.value = min(new_val, self.max_kl_coef)
+
+    def __call__(self) -> float:
+        return self.value
+
+
+def compute_kl_penalty(kl_vals: torch.Tensor, kl_coef: float, kl_penalty_cap: float):
+    """
+    Compute KL penalty with clipping.
+    Returns (clipped_penalty, raw_penalty, kl_mean).
+    """
+    kl_mean = kl_vals.mean()
+    raw_penalty = kl_coef * kl_mean
+    clipped_penalty = torch.clamp(raw_penalty, max=kl_penalty_cap)
+    return clipped_penalty, raw_penalty, kl_mean
+
+
+class EnhancedEntropyController:
+    def __init__(self, min_entropy: float = 0.5, max_entropy: float = 3.0,
+                 target_entropy: float = 1.5):
+        self.min_entropy = min_entropy
+        self.max_entropy = max_entropy
+        self.target_entropy = target_entropy
+        self.entropy_history: List[float] = []
+        self.entropy_weight = 0.01
+
+    def update_entropy_weight(self, current_entropy: float) -> float:
+        self.entropy_history.append(float(current_entropy))
+        if len(self.entropy_history) > 100:
+            self.entropy_history = self.entropy_history[-100:]
+        if len(self.entropy_history) >= 10:
+            avg_entropy = np.mean(self.entropy_history[-10:])
+            if avg_entropy < self.target_entropy * 0.8:
+                self.entropy_weight = min(0.05, self.entropy_weight * 1.1)
+            elif avg_entropy > self.target_entropy * 1.2:
+                self.entropy_weight = max(0.001, self.entropy_weight * 0.95)
+        return float(self.entropy_weight)
+
+    def adjust_for_seq_len(self, seq_len: int, base_entropy: float = 1.5):
+        seq_len = max(1, int(seq_len))
+        self.target_entropy = float(base_entropy * np.log1p(seq_len) / np.log1p(10))
+        self.target_entropy = float(np.clip(self.target_entropy, self.min_entropy, self.max_entropy))
+
+    def reset(self):
+        self.entropy_history.clear()
+        self.entropy_weight = 0.01
+
+
+class CurriculumManager:
+    """Symmetric curriculum: 10→15→20→25→20→15→10→..."""
+    def __init__(self, start_len: int = 10, max_len: int = 25,
+                 step_increase: int = 5, steps_per_level: int = 30):
+        self.start_len = start_len
+        self.max_len = max_len
+        self.step_increase = step_increase
+        self.steps_per_level = steps_per_level
+        self.current_max_len = start_len
+        self.step_counter = 0
+        self.direction = +1
+
+    def get_max_new_tokens(self) -> int:
+        return self.current_max_len
+
+    def step(self) -> int:
+        self.step_counter += 1
+        if self.step_counter % self.steps_per_level == 0:
+            if self.direction == +1:
+                if self.current_max_len < self.max_len:
+                    self.current_max_len += self.step_increase
+                else:
+                    self.direction = -1
+                    self.current_max_len -= self.step_increase
+            else:
+                if self.current_max_len > self.start_len:
+                    self.current_max_len -= self.step_increase
+                else:
+                    self.direction = +1
+                    self.current_max_len += self.step_increase
+            print(f"📈 Curriculum Update: max_new_tokens = {self.current_max_len}")
+        return self.current_max_len
+
+# ========================
+# HELPERS
+# ========================
+
+def normalize_rewards(rewards: torch.Tensor, seq_len: int, mode: str = "sqrt") -> torch.Tensor:
+    if seq_len <= 1 or mode == "none":
+        return rewards
+    if mode == "per_token":
+        return rewards / float(seq_len)
+    elif mode == "sqrt":
+        return rewards / float(np.sqrt(seq_len))
+    else:
+        raise ValueError(f"Unknown normalization mode: {mode}")
+
+
+def reset_controllers_on_phase_change(prev_len: Optional[int], new_len: int,
+                                      kl_controller: Optional[AdaptiveKLController] = None,
+                                      entropy_controller: Optional[EnhancedEntropyController] = None,
+                                      entropy_base: float = 1.5):
+    if prev_len is None or prev_len == new_len:
+        return
+    if kl_controller is not None:
+        kl_controller.reset()
+    if entropy_controller is not None:
+        entropy_controller.reset()
+        entropy_controller.adjust_for_seq_len(new_len, base_entropy=entropy_base)
+
+
+# ========================
+# PPO LOSS
+# ========================
+
+def compute_ppo_loss(old_log_probs: torch.Tensor, new_log_probs: torch.Tensor,
+                     rewards: torch.Tensor, clip_epsilon: float = 0.2,
+                     baseline: Optional[torch.Tensor] = None,
+                     seq_len: int = 1, reward_norm: str = "sqrt",
+                     adv_clip: Optional[float] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+    normed_rewards = normalize_rewards(rewards, seq_len, mode=reward_norm)
+    if baseline is not None:
+        advantage = normed_rewards - baseline.detach()
+    else:
+        advantage = normed_rewards
+    if adv_clip is not None:
+        advantage = torch.clamp(advantage, -float(adv_clip), float(adv_clip))
+    else:
+        default_clip = 2.0 * np.sqrt(max(1, seq_len))
+        advantage = torch.clamp(advantage, -default_clip, default_clip)
+    log_ratio = torch.clamp(new_log_probs - old_log_probs, -10.0, 10.0)
+    ratio = torch.exp(log_ratio)
+    adv_expanded = advantage.unsqueeze(1) if advantage.dim() == 1 else advantage
+    surr1 = ratio * adv_expanded
+    surr2 = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon) * adv_expanded
+    ppo_loss = -torch.min(surr1, surr2).sum(dim=1).mean()
+    return ppo_loss, advantage.detach()
+
+
+def compute_kl_divergence(old_action_probs: torch.Tensor, new_action_probs: torch.Tensor) -> torch.Tensor:
+    old_probs = old_action_probs.clamp_min(1e-12)
+    new_probs = new_action_probs.clamp_min(1e-12)
+    kl_per_step = (old_probs * (torch.log(old_probs) - torch.log(new_probs))).sum(dim=-1)
+    return kl_per_step.sum(dim=1)
+
+
+def compute_entropy_bonus(action_probs: torch.Tensor) -> torch.Tensor:
+    probs = action_probs.clamp_min(1e-12)
+    entropy_per_step = -(probs * torch.log(probs)).sum(dim=-1)
+    return entropy_per_step.sum(dim=1)
+
+# ========================
+# BATCH REWARD COMPUTATION
+# ========================
+
+def batch_compute_rewards_pareto(
+    selfies_list: List[str],
+    reward_mode: str = "mix",
+    reward_mix: float = 0.5,
+    pareto_controller: Optional[ParetoRewardController] = None
+) -> Dict[str, torch.Tensor]:
+    """
+    Drop-in replacement for batch_compute_rewards with Pareto support.
+    
+    Args:
+        selfies_list: List of SELFIES strings  
+        reward_mode: "chemq3", "sa", "mix", or "pareto" 
+        reward_mix: Weight for comprehensive rewards when mixing (0-1)
+        pareto_controller: ParetoRewardController instance for "pareto" mode
+        
+    Returns:
+        Dictionary containing reward tensors (same format as original)
+    """
+    batch_size = len(selfies_list)
+    
+    validity_vals = []
+    lipinski_vals = []
+    total_rewards = []
+    sa_rewards = []
+
+    # Compute all individual rewards
+    for selfies_str in selfies_list:
+        smiles = selfies_to_smiles(selfies_str)
+        
+        # Check validity comprehensively  
+        is_valid = (smiles is not None and 
+                   is_valid_smiles(smiles) and 
+                   passes_durrant_lab_filter(smiles))
+        
+        if reward_mode in ["chemq3", "mix", "pareto"]:
+            r = compute_comprehensive_reward(selfies_str)
+            validity_vals.append(r.get('validity', 0.0))
+            lipinski_vals.append(r.get('lipinski', 0.0))
+
+        if reward_mode in ["sa", "mix", "pareto"]:
+            sa = compute_sa_reward(selfies_str) if is_valid else 0.0
+            sa_rewards.append(sa)
+
+        # Store individual comprehensive reward for pareto mode
+        if reward_mode in ["chemq3", "pareto"]:
+            total_rewards.append(r.get('total', 0.0))
+        elif reward_mode == "sa":
+            total_rewards.append(sa)
+        elif reward_mode == "mix":
+            r_total = r.get("total", 0.0) if 'r' in locals() else 0.0
+            sa_val = sa if 'sa' in locals() else 0.0
+            mixed = reward_mix * r_total + (1.0 - reward_mix) * sa_val
+            total_rewards.append(mixed)
+
+    # Convert to tensors
+    result = {
+        "total_rewards": torch.tensor(total_rewards, dtype=torch.float32),
+    }
+    
+    if validity_vals:
+        result["validity_rewards"] = torch.tensor(validity_vals, dtype=torch.float32)
+    if lipinski_vals:
+        result["lipinski_rewards"] = torch.tensor(lipinski_vals, dtype=torch.float32)
+    if sa_rewards:
+        result["sa_rewards"] = torch.tensor(sa_rewards, dtype=torch.float32)
+    
+    # Compute diversity reward
+    valid_smiles = []
+    for selfies_str in selfies_list:
+        smiles = selfies_to_smiles(selfies_str)
+        if smiles and is_valid_smiles(smiles) and passes_durrant_lab_filter(smiles):
+            valid_smiles.append(smiles)
+    
+    diversity_score = len(set(valid_smiles)) / max(1, len(valid_smiles))
+    result["diversity_rewards"] = torch.full((batch_size,), diversity_score, dtype=torch.float32)
+    
+    # Apply Pareto mixing if requested
+    if reward_mode == "pareto" and pareto_controller is not None:
+        # Prepare objectives for controller update
+        batch_objectives = {
+            "total": result["total_rewards"],
+            "validity": result.get("validity_rewards", torch.zeros(batch_size)),
+            "diversity": result["diversity_rewards"]
+        }
+        
+        if "sa_rewards" in result:
+            batch_objectives["sa"] = result["sa_rewards"]
+        
+        # Update controller and get new weights
+        updated_weights = pareto_controller.update(batch_objectives)
+        
+        # Compute mixed reward using adaptive weights
+        mixed_reward = pareto_controller.get_mixed_reward(batch_objectives)
+        result["total_rewards"] = mixed_reward
+        
+        # Store weights for logging
+        result["pareto_weights"] = updated_weights
+
+    return result
+
+# Legacy
+def batch_compute_rewards(
+    selfies_list: List[str],
+    reward_mode: str = "chemq3",
+    reward_mix: float = 0.5
+) -> Dict[str, torch.Tensor]:
+    """
+    Compute rewards for a batch of SELFIES strings.
+    
+    Args:
+        selfies_list: List of SELFIES strings
+        reward_mode: "chemq3", "sa", or "mix"
+        reward_mix: Weight for chemq3 rewards when mixing (0-1)
+        
+    Returns:
+        Dictionary containing reward tensors
+    """
+    batch_size = len(selfies_list)
+    
+    validity_vals = []
+    lipinski_vals = []
+    total_rewards = []
+    sa_rewards = []
+
+    for selfies_str in selfies_list:
+        smiles = selfies_to_smiles(selfies_str)
+        
+        # Check validity comprehensively
+        is_valid = (smiles is not None and 
+                   is_valid_smiles(smiles) and 
+                   passes_durrant_lab_filter(smiles))
+        
+        if reward_mode == "chemq3":
+            r = compute_comprehensive_reward(selfies_str)
+            validity_vals.append(r.get('validity', 0.0))
+            lipinski_vals.append(r.get('lipinski', 0.0))
+            total_rewards.append(r.get('total', 0.0))
+
+        elif reward_mode == "sa":
+            sa = compute_sa_reward(selfies_str) if is_valid else 0.0
+            sa_rewards.append(sa)
+            total_rewards.append(sa)
+
+        elif reward_mode == "mix":
+            r = compute_comprehensive_reward(selfies_str)
+            sa = compute_sa_reward(selfies_str) if is_valid else 0.0
+            mixed = reward_mix * r.get("total", 0.0) + (1.0 - reward_mix) * sa
+            
+            total_rewards.append(mixed)
+            sa_rewards.append(sa)
+            validity_vals.append(r.get('validity', 0.0))
+            lipinski_vals.append(r.get('lipinski', 0.0))
+
+        else:
+            # Unknown mode -> default to zero reward
+            total_rewards.append(0.0)
+            validity_vals.append(0.0)
+            lipinski_vals.append(0.0)
+
+    # Convert to tensors
+    result = {
+        "total_rewards": torch.tensor(total_rewards, dtype=torch.float32),
+    }
+    
+    if validity_vals:
+        result["validity_rewards"] = torch.tensor(validity_vals, dtype=torch.float32)
+    if lipinski_vals:
+        result["lipinski_rewards"] = torch.tensor(lipinski_vals, dtype=torch.float32)
+    if sa_rewards:
+        result["sa_rewards"] = torch.tensor(sa_rewards, dtype=torch.float32)
+        
+    return result
+
+# ========================
+# TRAINING METRICS
+# ========================
+
+def compute_training_metrics(
+    rewards: Dict[str, torch.Tensor],
+    selfies_list: List[str],
+    loss_dict: Dict[str, float]
+) -> Dict[str, float]:
+    """
+    Compute comprehensive training metrics.
+    
+    Args:
+        rewards: Dictionary of reward tensors
+        selfies_list: List of generated SELFIES
+        loss_dict: Dictionary containing loss components
+        
+    Returns:
+        Dictionary of computed metrics
+    """
+    metrics = {}
+    
+    # Basic reward metrics
+    if "total_rewards" in rewards:
+        metrics["avg_reward"] = float(rewards["total_rewards"].mean())
+        metrics["max_reward"] = float(rewards["total_rewards"].max())
+        metrics["min_reward"] = float(rewards["total_rewards"].min())
+        metrics["reward_std"] = float(rewards["total_rewards"].std())
+    
+    if "validity_rewards" in rewards:
+        metrics["validity_rate"] = float(rewards["validity_rewards"].mean())
+    
+    if "lipinski_rewards" in rewards:
+        metrics["lipinski_score"] = float(rewards["lipinski_rewards"].mean())
+    
+    if "sa_rewards" in rewards:
+        metrics["sa_score"] = float(rewards["sa_rewards"].mean())
+    
+    # Molecular diversity metrics
+    valid_smiles = []
+    for selfies_str in selfies_list:
+        smiles = selfies_to_smiles(selfies_str)
+        if smiles and is_valid_smiles(smiles) and passes_durrant_lab_filter(smiles):
+            valid_smiles.append(smiles)
+    
+    metrics["num_valid"] = len(valid_smiles)
+    metrics["num_unique"] = len(set(valid_smiles))
+    metrics["diversity_ratio"] = len(set(valid_smiles)) / max(1, len(valid_smiles))
+    
+    # Add loss components
+    metrics.update(loss_dict)
+    
+
+    return metrics

ChemQ3MTP/trainer.py

@@ -0,0 +1,72 @@
+# trainer.py
+from transformers import Trainer, TrainingArguments
+from torch.utils.data import DataLoader
+import torch
+import torch.nn.functional as F
+
+class MTPTrainer(Trainer):
+    """
+    Custom trainer for Multi-Token Prediction training.
+    """
+    def __init__(self, model, args=None, train_dataset=None, eval_dataset=None, **kwargs):
+        super().__init__(
+            model=model,
+            args=args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            **kwargs
+        )
+        self.use_mtp_training = True
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """
+        Compute loss during training - handles both MTP and standard LM training.
+        """
+        labels = inputs.get("labels")
+        outputs = model(**inputs, labels=labels, use_mtp_training=self.use_mtp_training)
+        loss = outputs.loss
+        
+        return (loss, outputs) if return_outputs else loss
+
+    def train_step_with_mtp(self, model, inputs):
+        """
+        Specialized training step for MTP training.
+        """
+        model.set_mtp_training(True)
+        return self.training_step(model, inputs)
+
+    def train_step_with_lm(self, model, inputs):
+        """
+        Standard language modeling training step.
+        """
+        model.set_mtp_training(False)
+        return self.training_step(model, inputs)
+
+
+class RLTrainer:
+    """
+    Separate trainer class for Reinforcement Learning training.
+    This can use the generate_with_logprobs method from your model.
+    """
+    def __init__(self, model, tokenizer, rl_config=None):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.rl_config = rl_config or {}
+        
+    def rl_training_step(self, input_ids, old_log_probs, old_action_probs, **kwargs):
+        """
+        Perform an RL training step using the model's generate_with_logprobs method
+        and the reward functions from rl_utils.
+        """
+        # Import RL utilities
+        from .rl_utils import (
+            batch_compute_rewards, 
+            compute_ppo_loss, 
+            compute_kl_divergence,
+            compute_entropy_bonus,
+            AdaptiveKLController
+        )
+        
+        # This would call the generate_with_logprobs method from your model
+        # and then compute RL-specific losses
+        pass

ChemQ3MTP/training_utils.py

@@ -0,0 +1,157 @@
+# training_utils.py
+import torch
+import torch.nn as nn
+from transformers import Trainer, TrainingArguments
+from datasets import load_dataset, Dataset
+from ranger21 import Ranger21
+from torch.optim.lr_scheduler import LambdaLR
+import os
+from typing import Dict, Any, Tuple
+from transformers import TrainerCallback
+
+class EnhancedDataCollator:
+    def __init__(self, tokenizer, pad_to_multiple_of=8):
+        self.tokenizer = tokenizer
+        self.pad_to_multiple_of = pad_to_multiple_of
+
+    def __call__(self, features):
+        max_length = max(len(f["input_ids"]) for f in features)
+        if self.pad_to_multiple_of:
+            max_length = ((max_length + self.pad_to_multiple_of - 1) // self.pad_to_multiple_of) * self.pad_to_multiple_of
+        
+        batch = {"input_ids": [], "attention_mask": [], "labels": []}
+        for feature in features:
+            input_ids = feature["input_ids"]
+            attention_mask = feature["attention_mask"]
+            labels = feature["labels"]
+            padding_length = max_length - len(input_ids)
+            padded_input_ids = input_ids + [self.tokenizer.pad_token_id] * padding_length
+            padded_attention_mask = attention_mask + [0] * padding_length
+            padded_labels = labels + [-100] * padding_length
+            batch["input_ids"].append(padded_input_ids)
+            batch["attention_mask"].append(padded_attention_mask)
+            batch["labels"].append(padded_labels)
+        
+        batch = {k: torch.tensor(v, dtype=torch.long) for k, v in batch.items()}
+        return batch
+
+def tokenize_function(examples, tokenizer, max_length):
+    smiles_list = examples['SELFIES']
+    batch_results = {"input_ids": [], "attention_mask": [], "labels": []}
+    for smiles in smiles_list:
+        tokenized = tokenizer(
+            smiles,
+            truncation=True,
+            padding=False,
+            max_length=max_length,
+            return_tensors=None,
+            add_special_tokens=True
+        )
+        input_ids = tokenized["input_ids"]
+        attention_mask = tokenized["attention_mask"]
+        labels = input_ids.copy()
+        batch_results["input_ids"].append(input_ids)
+        batch_results["attention_mask"].append(attention_mask)
+        batch_results["labels"].append(labels)
+    return batch_results
+
+def load_and_tokenize_dataset(data_path: str, tokenizer, max_length: int, tokenize_batch_size: int, 
+                              train_split_ratio: float, val_split_ratio: float, test_split_ratio: float):
+    dataset = load_dataset('csv', data_files=data_path, split='train')
+    dataset = dataset.shuffle(seed=42)
+
+    total_lines = len(dataset)
+    test_size = int(test_split_ratio * total_lines)
+    val_size = int(val_split_ratio * total_lines)
+    train_size = total_lines - test_size - val_size
+
+    train_dataset = dataset.select(range(0, train_size))
+    val_dataset = dataset.select(range(train_size, train_size + val_size))
+    test_dataset = dataset.select(range(train_size + val_size, total_lines))
+
+    def tokenize_train(examples):
+        return tokenize_function(examples, tokenizer, max_length)
+    
+    def tokenize_val(examples):
+        return tokenize_function(examples, tokenizer, max_length)
+
+    train_dataset = train_dataset.map(
+        tokenize_train, 
+        batched=True, 
+        batch_size=tokenize_batch_size, 
+        remove_columns=["SELFIES"],
+        desc="Tokenizing train"
+    )
+    val_dataset = val_dataset.map(
+        tokenize_val, 
+        batched=True, 
+        batch_size=tokenize_batch_size, 
+        remove_columns=["SELFIES"],
+        desc="Tokenizing val"
+    )
+
+    return train_dataset, val_dataset, test_dataset
+
+def create_enhanced_optimizer(model_params, lr, weight_decay, num_epochs, train_dataset_len, batch_size):
+    num_batches_per_epoch = train_dataset_len // batch_size
+    optimizer_params = {
+        'lr': lr,
+        'weight_decay': weight_decay,
+        'use_adabelief': True,
+        'use_madgrad': True,
+        'using_gc': True,
+        'warmdown_active': True,
+        'num_epochs': num_epochs,
+        'num_batches_per_epoch': num_batches_per_epoch,
+        'use_warmup': True,
+        'use_cheb': False
+    }
+    return Ranger21(model_params, **optimizer_params)
+
+class EnhancedCustomTrainer(Trainer):
+    def create_optimizer(self):
+        self.optimizer = create_enhanced_optimizer(
+            self.model.parameters(),
+            self.args.learning_rate,
+            self.args.weight_decay,
+            self.args.num_train_epochs,
+            len(self.train_dataset),
+            self.args.per_device_train_batch_size
+        )
+        return self.optimizer
+
+    def create_scheduler(self, num_training_steps, optimizer=None):
+        self.lr_scheduler = LambdaLR(optimizer or self.optimizer, lr_lambda=lambda step: 1.0)
+        return self.lr_scheduler
+
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+        outputs = model(**inputs)
+        loss = outputs.loss
+        return (loss, outputs) if return_outputs else loss
+
+def setup_training_args(output_dir: str, total_steps: int, batch_size: int, gradient_accumulation_steps: int, 
+                        steps_per_epoch: int, include_for_metrics: list):
+    return TrainingArguments(
+        output_dir=output_dir,
+        max_steps=total_steps,
+        per_device_train_batch_size=batch_size,
+        per_device_eval_batch_size=batch_size,
+        gradient_accumulation_steps=gradient_accumulation_steps,
+        logging_dir='./logs',
+        logging_strategy="steps",
+        logging_steps=max(1, steps_per_epoch // 4),
+        eval_strategy="steps",
+        eval_steps=max(1, steps_per_epoch // 4),
+        save_strategy="steps",
+        save_steps=steps_per_epoch,
+        save_total_limit=1,
+        dataloader_num_workers=0,
+        dataloader_pin_memory=False,
+        remove_unused_columns=False,
+        prediction_loss_only=False,
+        fp16=torch.cuda.is_available(),
+        gradient_checkpointing=True,
+        dataloader_drop_last=True,
+        report_to=None,
+        include_for_metrics=include_for_metrics,
+    )

FastChemTokenizerHF.py

@@ -0,0 +1,539 @@
+import torch
+import json
+import os
+from typing import List, Union, Optional, Tuple, Dict, Any
+from functools import lru_cache
+from collections.abc import Mapping
+
+
+# ------------------------------
+# BatchEncoding
+# ------------------------------
+class BatchEncoding(dict, Mapping):
+    """Minimal BatchEncoding compatible wrapper."""
+
+    def __init__(self, data: dict, tensor_type: Optional[str] = None):
+        data = {} if data is None else {k: v for k, v in data.items()}
+        super().__init__(data)
+        self.data = data
+        self.tensor_type = tensor_type
+        for k, v in data.items():
+            setattr(self, k, v)
+
+    def __getitem__(self, key): return self.data[key]
+    def __iter__(self): return iter(self.data)
+    def __len__(self): return len(self.data)
+    def keys(self): return self.data.keys()
+    def values(self): return self.data.values()
+    def items(self): return self.data.items()
+    def get(self, key, default=None): return self.data.get(key, default)
+
+    def to(self, device):
+        if self.tensor_type in ("pt", "torch"):
+            for k, v in list(self.data.items()):
+                if torch.is_tensor(v):
+                    self.data[k] = v.to(device)
+                    setattr(self, k, self.data[k])
+        return self
+
+    def cpu(self): return self.to("cpu")
+    def cuda(self): return self.to("cuda")
+    def detach(self):
+        if self.tensor_type in ("pt", "torch"):
+            for k, v in list(self.data.items()):
+                if torch.is_tensor(v):
+                    self.data[k] = v.detach()
+                    setattr(self, k, self.data[k])
+        return self
+
+    def __repr__(self):
+        keys = ", ".join(list(self.data.keys())[:10])
+        return f"BatchEncoding(keys=[{keys}], tensor_type={self.tensor_type})"
+
+
+# ------------------------------
+# Base class
+# ------------------------------
+class PreTrainedTokenizerBase:
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            if key.endswith('_token'):
+                setattr(self, f"_{key}", value)
+                setattr(self, f"{key}_id", None)
+        self.model_max_length = kwargs.get('model_max_length', 512)
+        self.padding_side = kwargs.get('padding_side', 'right')
+        self.truncation_side = kwargs.get('truncation_side', 'right')
+        self.chat_template = kwargs.get('chat_template')
+
+
+# ------------------------------
+# Trie node
+# ------------------------------
+class TrieNode:
+    __slots__ = ['children', 'token_id']
+    def __init__(self):
+        self.children = {}
+        self.token_id = None
+
+
+# ------------------------------
+# FastChemTokenizer
+# ------------------------------
+
+class FastChemTokenizer(PreTrainedTokenizerBase):
+    def __init__(self, token_to_id=None, vocab_file=None, **kwargs):
+        if vocab_file is not None:
+            with open(vocab_file, "r", encoding="utf-8") as f:
+                token_to_id = json.load(f)
+                token_to_id = {str(k): int(v) for k, v in token_to_id.items()}
+
+        self.token_to_id = token_to_id
+        self.id_to_token = {v: k for k, v in token_to_id.items()}
+
+        # Build trie
+        self.trie_root = self._build_trie(self.token_to_id)
+
+        # ✅ Call parent (sets token *strings*, may reset *_id to None)
+        super().__init__(
+            bos_token="<s>",
+            eos_token="</s>",
+            unk_token="<unk>",
+            pad_token="<pad>",
+            mask_token="<mask>",
+            model_max_length=kwargs.get("model_max_length", 512),
+            padding_side=kwargs.get("padding_side", "right"),
+            truncation_side=kwargs.get("truncation_side", "right"),
+            **kwargs,
+        )
+
+        # ✅ Re-map token strings → IDs from vocab
+        self.bos_token_id  = self.token_to_id.get("<s>", 0)
+        self.eos_token_id  = self.token_to_id.get("</s>", 1)
+        self.pad_token_id  = self.token_to_id.get("<pad>", 2)
+        self.unk_token_id  = self.token_to_id.get("<unk>", 3)
+        self.mask_token_id = self.token_to_id.get("<mask>", 4)
+
+        # Ensure reverse mapping always valid
+        self.id_to_token[self.bos_token_id]  = "<s>"
+        self.id_to_token[self.eos_token_id]  = "</s>"
+        self.id_to_token[self.pad_token_id]  = "<pad>"
+        self.id_to_token[self.unk_token_id]  = "<unk>"
+        self.id_to_token[self.mask_token_id] = "<mask>"
+
+        # Debug
+        print("✅ Special tokens bound:",
+              self.bos_token_id, self.eos_token_id, self.pad_token_id,
+              self.unk_token_id, self.mask_token_id)
+        
+        # ✅ Ensure token *strings* also exist (for decode fallback)
+        self.bos_token = "<s>"
+        self.eos_token = "</s>"
+        self.pad_token = "<pad>"
+        self.unk_token = "<unk>"
+        self.mask_token = "<mask>"
+
+
+    def _build_trie(self, token_to_id):
+        root = TrieNode()
+        for token, tid in token_to_id.items():
+            node = root
+            for char in token:
+                if char not in node.children:
+                    node.children[char] = TrieNode()
+                node = node.children[char]
+            node.token_id = tid
+        return root
+
+    @property
+    def vocab_size(self): return len(self.token_to_id)
+    def __len__(self): return len(self.token_to_id)
+    def get_vocab(self) -> Dict[str, int]: return self.token_to_id.copy()
+
+    @lru_cache(maxsize=10000)
+    def _cached_encode_str(self, s: str) -> Tuple[int, ...]:
+        return tuple(self._encode_core(s))
+
+    def _encode_core(self, text: str) -> List[int]:
+        tokens, result_ids = text, []
+        i, n = 0, len(tokens)
+        while i < n:
+            node, j = self.trie_root, i
+            last_match_id, last_match_end = None, i
+            while j < n and tokens[j] in node.children:
+                node = node.children[tokens[j]]
+                j += 1
+                if node.token_id is not None:
+                    last_match_id, last_match_end = node.token_id, j
+            if last_match_id is not None:
+                result_ids.append(last_match_id)
+                i = last_match_end
+            else:
+                tid = self.token_to_id.get(tokens[i], self.unk_token_id)
+                result_ids.append(tid)
+                i += 1
+        return result_ids
+
+    # ------------------------------
+    # Converters
+    # ------------------------------
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.token_to_id.get(token, self.unk_token_id)
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.id_to_token.get(index, self.unk_token)
+
+    def convert_tokens_to_ids(self, tokens: Union[str, List[str]]):
+        if isinstance(tokens, str): return self._convert_token_to_id(tokens)
+        return [self._convert_token_to_id(tok) for tok in tokens]
+
+    def convert_ids_to_tokens(self, ids: Union[int, List[int]]):
+        if isinstance(ids, int): return self._convert_id_to_token(ids)
+        return [self._convert_id_to_token(i) for i in ids]
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str: return "".join(tokens)
+
+    # ------------------------------
+    # Encoding / Decoding
+    # ------------------------------
+        # ------------------------------
+    # Convenience wrappers
+    # ------------------------------
+    def encode(
+        self,
+        text: str,
+        text_pair: Optional[str] = None,
+        add_special_tokens: bool = True,
+        padding: bool = False,
+        truncation: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+    ) -> List[int]:
+        encoded = self.encode_plus(
+            text=text,
+            text_pair=text_pair,
+            add_special_tokens=add_special_tokens,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+        )
+        input_ids = encoded["input_ids"]
+        if isinstance(input_ids, torch.Tensor):
+            if input_ids.dim() > 1:
+                input_ids = input_ids.squeeze(0)
+            input_ids = input_ids.tolist()
+        return input_ids
+
+    def __call__(
+        self, 
+        text: Union[str, List[str]], 
+        text_pair: Optional[Union[str, List[str]]] = None, 
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        is_split_into_words: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, Any]] = None,
+        return_token_type_ids: Optional[bool] = None,
+        return_attention_mask: Optional[bool] = None,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        **kwargs
+    ) -> BatchEncoding:
+        """HuggingFace-compatible: one string → encode_plus, list → batch_encode_plus"""
+        if return_token_type_ids is None:
+            return_token_type_ids = True
+        if return_attention_mask is None:
+            return_attention_mask = True
+
+        if isinstance(text, list):
+            if text_pair is not None:
+                batch = [(t, p) for t, p in zip(text, text_pair)]
+            else:
+                batch = text
+            return self.batch_encode_plus(
+                batch,
+                add_special_tokens=add_special_tokens,
+                padding=padding,
+                truncation=truncation,
+                max_length=max_length,
+                stride=stride,
+                is_split_into_words=is_split_into_words,
+                pad_to_multiple_of=pad_to_multiple_of,
+                return_tensors=return_tensors,
+                return_token_type_ids=return_token_type_ids,
+                return_attention_mask=return_attention_mask,
+                return_overflowing_tokens=return_overflowing_tokens,
+                return_special_tokens_mask=return_special_tokens_mask,
+                return_offsets_mapping=return_offsets_mapping,
+                return_length=return_length,
+                verbose=verbose,
+                **kwargs
+            )
+        else:
+            return self.encode_plus(
+                text=text,
+                text_pair=text_pair,
+                add_special_tokens=add_special_tokens,
+                padding=padding,
+                truncation=truncation,
+                max_length=max_length,
+                stride=stride,
+                is_split_into_words=is_split_into_words,
+                pad_to_multiple_of=pad_to_multiple_of,
+                return_tensors=return_tensors,
+                return_token_type_ids=return_token_type_ids,
+                return_attention_mask=return_attention_mask,
+                return_overflowing_tokens=return_overflowing_tokens,
+                return_special_tokens_mask=return_special_tokens_mask,
+                return_offsets_mapping=return_offsets_mapping,
+                return_length=return_length,
+                verbose=verbose,
+                **kwargs
+            )
+
+    def encode_plus(
+        self,
+        text: str,
+        text_pair: Optional[str] = None,
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        is_split_into_words: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, Any]] = None,
+        return_token_type_ids: Optional[bool] = True,
+        return_attention_mask: Optional[bool] = True,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        **kwargs
+    ) -> BatchEncoding:
+        if max_length is None: max_length = self.model_max_length
+        ids_a = list(self._cached_encode_str(text.strip()))
+        ids_b = list(self._cached_encode_str(text_pair.strip())) if text_pair else None
+
+        input_ids, token_type_ids = [], []
+        if add_special_tokens:
+            input_ids.append(self.bos_token_id); token_type_ids.append(0)
+            input_ids.extend(ids_a); token_type_ids.extend([0] * len(ids_a))
+            input_ids.append(self.eos_token_id); token_type_ids.append(0)
+            if ids_b is not None:
+                input_ids.extend(ids_b); token_type_ids.extend([1] * len(ids_b))
+                input_ids.append(self.eos_token_id); token_type_ids.append(1)
+        else:
+            input_ids = ids_a.copy(); token_type_ids = [0] * len(input_ids)
+            if ids_b is not None:
+                input_ids.extend(ids_b); token_type_ids.extend([1] * len(ids_b))
+
+        if truncation and len(input_ids) > max_length:
+            input_ids, token_type_ids = input_ids[:max_length], token_type_ids[:max_length]
+
+        encoded_dict = {"input_ids": input_ids}
+        if return_attention_mask:
+            if padding == True or padding == "max_length":
+                pad_len = max_length - len(input_ids)
+                if pad_len > 0:
+                    if self.padding_side == "right":
+                        input_ids.extend([self.pad_token_id] * pad_len)
+                        token_type_ids.extend([0] * pad_len)
+                    else:
+                        input_ids = [self.pad_token_id] * pad_len + input_ids
+                        token_type_ids = [0] * pad_len + token_type_ids
+            attention_mask = [0 if tid == self.pad_token_id else 1 for tid in input_ids]
+            encoded_dict["attention_mask"] = attention_mask
+        if return_token_type_ids: encoded_dict["token_type_ids"] = token_type_ids
+        if return_special_tokens_mask:
+            encoded_dict["special_tokens_mask"] = [
+                1 if tid in {self.bos_token_id, self.eos_token_id, self.pad_token_id, self.mask_token_id} else 0
+                for tid in input_ids
+            ]
+        if return_length:
+            encoded_dict["length"] = len([tid for tid in input_ids if tid != self.pad_token_id])
+
+        if return_tensors in ["pt", "torch"]:
+            out = {}
+            for k, v in encoded_dict.items():
+                if isinstance(v, list):
+                    tensor = torch.tensor(
+                        [self.unk_token_id if x is None else int(x) for x in v], dtype=torch.long
+                    ).unsqueeze(0)
+                    out[k] = tensor
+                else:
+                    out[k] = v
+            return BatchEncoding(out, tensor_type=return_tensors)
+        return BatchEncoding(encoded_dict, tensor_type=None)
+
+    def batch_encode_plus(
+        self,
+        batch_text_or_text_pairs: List[Union[str, Tuple[str, str]]],
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        is_split_into_words: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, Any]] = None,
+        return_token_type_ids: Optional[bool] = True,
+        return_attention_mask: Optional[bool] = True,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        **kwargs
+    ) -> BatchEncoding:
+        if padding is True: padding = "longest"
+        if padding == "max_length" and max_length is None: max_length = self.model_max_length
+
+        all_input_ids, all_token_type_ids, all_attention_masks = [], [], []
+        all_special_masks, all_lengths = [], []
+        for item in batch_text_or_text_pairs:
+            t, tp = item if isinstance(item, tuple) else (item, None)
+            enc = self.encode_plus(
+                text=t, text_pair=tp, add_special_tokens=add_special_tokens,
+                padding=False, truncation=truncation, max_length=max_length,
+                return_tensors=None, return_token_type_ids=return_token_type_ids,
+                return_attention_mask=return_attention_mask,
+                return_special_tokens_mask=return_special_tokens_mask,
+                return_length=return_length, **kwargs
+            )
+            ids, tt, am = enc["input_ids"], enc.get("token_type_ids", [0]*len(enc["input_ids"])), enc.get("attention_mask",[1]*len(enc["input_ids"]))
+            sm, ln = enc.get("special_tokens_mask",[0]*len(ids)), enc.get("length", len([x for x in ids if x != self.pad_token_id]))
+            all_input_ids.append(ids); all_token_type_ids.append(tt); all_attention_masks.append(am)
+            all_special_masks.append(sm); all_lengths.append(ln)
+
+        pad_to = max(len(x) for x in all_input_ids) if padding == "longest" else (max_length if padding == "max_length" else None)
+        batched = {
+            "input_ids": all_input_ids,
+            "token_type_ids": all_token_type_ids if return_token_type_ids else None,
+            "attention_mask": all_attention_masks if return_attention_mask else None,
+            "special_tokens_mask": all_special_masks if return_special_tokens_mask else None,
+            "length": all_lengths if return_length else None,
+        }
+        if pad_to is not None:
+            for key in ["input_ids","token_type_ids","attention_mask","special_tokens_mask"]:
+                if batched.get(key) is None: continue
+                padded = []
+                for seq in batched[key]:
+                    pad_len = pad_to - len(seq)
+                    pad_val = self.pad_token_id if key=="input_ids" else 0
+                    if pad_len > 0:
+                        seq = seq+[pad_val]*pad_len if self.padding_side=="right" else [pad_val]*pad_len+seq
+                    padded.append(seq)
+                batched[key] = padded
+
+        if return_tensors in ["pt", "torch"]:
+            def to_tensor(lst, pad_val=0):
+                return torch.tensor([[self.unk_token_id if x is None else int(x) for x in row] for row in lst], dtype=torch.long)
+            out = {}
+            if batched.get("input_ids") is not None: out["input_ids"] = to_tensor(batched["input_ids"], self.pad_token_id)
+            if batched.get("attention_mask") is not None: out["attention_mask"] = to_tensor(batched["attention_mask"],0)
+            if batched.get("token_type_ids") is not None: out["token_type_ids"] = to_tensor(batched["token_type_ids"],0)
+            if batched.get("special_tokens_mask") is not None: out["special_tokens_mask"] = to_tensor(batched["special_tokens_mask"],0)
+            if return_length and batched.get("length") is not None: out["length"] = torch.tensor([int(x) for x in batched["length"]], dtype=torch.long)
+            return BatchEncoding(out, tensor_type=return_tensors)
+        return BatchEncoding({k:v for k,v in batched.items() if v is not None}, tensor_type=None)
+
+    # ------------------------------
+    # Decoding
+    # ------------------------------
+    def decode(self, token_ids, skip_special_tokens=False, **kwargs):
+        if isinstance(token_ids, torch.Tensor): token_ids = token_ids.tolist()
+        special_ids = {self.bos_token_id,self.eos_token_id,self.pad_token_id,self.mask_token_id} if skip_special_tokens else set()
+        tokens = [self.id_to_token.get(tid,self.unk_token) for tid in token_ids if tid not in special_ids]
+        return "".join(tokens)
+
+    def batch_decode(self, sequences, skip_special_tokens=False, **kwargs):
+        if isinstance(sequences, torch.Tensor): sequences = sequences.tolist()
+        return [self.decode(seq, skip_special_tokens=skip_special_tokens, **kwargs) for seq in sequences]
+
+    def decode_with_trace(self, token_ids: List[int]):
+        print(f"\n🔍 Decoding {len(token_ids)} tokens:")
+        for i, tid in enumerate(token_ids):
+            token = self.id_to_token.get(tid, self.unk_token)
+            tid_str = "None" if tid is None else f"{tid:5d}"
+            print(f"  [{i:03d}] ID={tid_str} → '{token}'")
+
+    # ------------------------------
+    # Save / Load
+    # ------------------------------
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory): os.makedirs(save_directory)
+        vocab_file = os.path.join(save_directory,(filename_prefix+"-" if filename_prefix else "")+"vocab.json")
+        with open(vocab_file,"w",encoding="utf-8") as f: json.dump(self.token_to_id,f,ensure_ascii=False,indent=2)
+        return (vocab_file,)
+
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], filename_prefix: Optional[str]=None, **kwargs):
+        if not os.path.exists(save_directory): os.makedirs(save_directory)
+        self.save_vocabulary(save_directory, filename_prefix)
+        config_file = os.path.join(save_directory,"tokenizer_config.json")
+        with open(config_file,"w",encoding="utf-8") as f:
+            json.dump({
+                "tokenizer_class": self.__class__.__name__,
+                "model_max_length": self.model_max_length,
+                "padding_side": self.padding_side,
+                "truncation_side": self.truncation_side,
+                "special_tokens": {
+                    "bos_token": self.bos_token,
+                    "eos_token": self.eos_token,
+                    "pad_token": self.pad_token,
+                    "unk_token": self.unk_token,
+                    "mask_token": self.mask_token,
+                }
+            },f,ensure_ascii=False,indent=2)
+        return (save_directory,)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        if os.path.isdir(pretrained_model_name_or_path):
+            vocab_file = os.path.join(pretrained_model_name_or_path,"vocab.json")
+            config_file = os.path.join(pretrained_model_name_or_path,"tokenizer_config.json")
+            config = {}
+            if os.path.exists(config_file):
+                with open(config_file,"r",encoding="utf-8") as f: config=json.load(f)
+            return cls(vocab_file=vocab_file, **{**config,**kwargs})
+        else:
+            raise NotImplementedError("Loading from Hub not implemented yet")
+
+
+# ------------------------------
+# SELFIES variant
+# ------------------------------
+class FastChemTokenizerSelfies(FastChemTokenizer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)  # ✅ ensures BOS/EOS etc. are set
+
+    """SELFIES variant that handles whitespace-separated tokens."""
+
+    def _encode_core(self, text: str) -> List[int]:
+        result_ids, i, n = [], 0, len(text)
+        while i < n:
+            if text[i].isspace(): i += 1; continue
+            node, j = self.trie_root, i
+            last_match_id, last_match_end = None, i
+            while j < n and text[j] in node.children:
+                node = node.children[text[j]]; j += 1
+                if node.token_id is not None:
+                    last_match_id, last_match_end = node.token_id, j
+            if last_match_id is not None:
+                result_ids.append(last_match_id); i = last_match_end
+            else:
+                result_ids.append(self.token_to_id.get(text[i], self.unk_token_id)); i += 1
+        return result_ids
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str: return " ".join(tokens)
+    def decode(self, token_ids, skip_special_tokens=False, **kwargs):
+        if isinstance(token_ids, torch.Tensor): token_ids = token_ids.tolist()
+        special_ids = {self.bos_token_id,self.eos_token_id,self.pad_token_id,self.mask_token_id} if skip_special_tokens else set()
+        tokens = [self.id_to_token.get(tid,self.unk_token) for tid in token_ids if tid not in special_ids]
+        return " ".join(tokens)