Upload 7 files

ChemQ3MTP.py

@@ -0,0 +1,753 @@
+# ========================
+#  ChemQ3-MTP
+#  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 Qwen3Config, Qwen3ForCausalLM, AutoTokenizer
+from rdkit import Chem
+from rdkit.Chem import Descriptors, Lipinski
+import selfies as sf
+from rdkit import RDLogger
+RDLogger.DisableLog('rdApp.*')      # suppress all SMILES parse messages
+import json
+from typing import List, Union, Optional, Tuple
+from transformers.tokenization_utils_base import BatchEncoding
+from FastChemTokenizer import FastChemTokenizerSelfies
+import numpy as np
+from collections import Counter
+from rdkit.Chem import Descriptors, Lipinski, rdMolDescriptors
+
+# ========================
+# UTILS: SELFIES -> SMILES -> VALIDITY & LIPINSKI
+# ========================
+
+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
+
+# SA Classifier
+from transformers import pipeline
+
+# Optional: lazy load so we don’t reload every time
+_sa_classifier = None
+def get_sa_classifier():
+    global _sa_classifier
+    if _sa_classifier is None:
+        _sa_classifier = pipeline("text-classification", model="gbyuvd/synthaccess-chemselfies")
+    return _sa_classifier
+
+
+def compute_sa_reward(selfies_str: str) -> float:
+    """Reward molecules with easy synthetic accessibility (SA)."""
+    try:
+        classifier = get_sa_classifier()
+        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
+
+# ==========================
+#  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 = []
+
+        # MW
+        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)
+
+        # Donors
+        scores.append(1.0 if hbd <= 5 else max(0.0, 1.0 - 0.2 * (hbd - 5)))
+        # 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
+
+
+# ==========================
+#  Unified Reward
+# ==========================
+def compute_comprehensive_reward(selfies_str: str) -> dict[str, float]:
+    smiles = selfies_to_smiles(selfies_str)
+    mol = Chem.MolFromSmiles(smiles) if smiles else None
+
+    rewards = {
+        "validity": 1.0 if mol is not None 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 rewards["validity"] == 0:
+        rewards["total"] = 0.0
+    else:
+        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 compute_lipinski_reward(mol) -> float:
+    if mol is None:
+        return 0.0
+    try:
+        mw = Descriptors.MolWt(mol)
+        logp = Descriptors.MolLogP(mol)
+        hbd = Lipinski.NumHDonors(mol)
+        hba = Lipinski.NumHAcceptors(mol)
+        rules = [250 < mw <= 500, logp <= 5, hbd <= 5, hba <= 10]   # we dont want too small of fragments
+        return sum(rules) / 4.0
+    except:
+        return 0.0
+
+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:
+        return 0.0
+    mol = Chem.MolFromSmiles(smiles)
+    return compute_lipinski_reward(mol)
+
+class MTPHead(nn.Module):
+    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):
+    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 ChemQ3MTP(Qwen3ForCausalLM):
+    def __init__(self, config, num_future_tokens: int = 3):
+        super().__init__(config)
+        self.mtp_head = MTPHead(config.hidden_size, config.vocab_size, num_future_tokens)
+        self.horizon_loss = HorizonLoss(num_future_tokens=num_future_tokens)
+        self.use_mtp_training = True
+        self.post_init()
+        self.entropy_controller = EnhancedEntropyController(
+            min_entropy=0.5,
+            max_entropy=3.0,
+            target_entropy=1.5,
+            adaptation_rate=0.01,
+            )
+
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,  
+        **kwargs
+    ):
+        # Default mask if not provided
+        if attention_mask is None and input_ids is not None:
+            attention_mask = (input_ids != self.config.pad_token_id).long()
+
+        # Respect caller settings, only set defaults if missing
+        kwargs.setdefault("output_hidden_states", True)
+        kwargs.setdefault("return_dict", True)
+
+        outputs = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=None,  
+            **kwargs
+        )
+
+        hidden_states = outputs.hidden_states[-1]
+        lm_logits = outputs.logits
+        loss = None
+
+        if self.training and self.use_mtp_training and labels is not None:  # labels, not kwargs
+            mtp_outputs = self.mtp_head(hidden_states)
+            horizon_loss_dict = self.horizon_loss(mtp_outputs, input_ids, attention_mask)
+
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()  # labels, not kwargs["labels"]
+
+            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')
+
+            loss = 0.7 * horizon_loss_dict['loss'] + 0.3 * causal_lm_loss
+
+        elif labels is not None:  #  labels, not kwargs.get("labels")
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()  #  labels, not kwargs["labels"]
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1)
+            )
+
+        from transformers.modeling_outputs import CausalLMOutputWithPast
+        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):
+        self.use_mtp_training = use_mtp
+
+    # ================
+    # RL SAMPLING + PPO
+    # ================
+
+    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,   # allow passing explicitly
+    ) -> Tuple[List[str], torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        self.eval()
+        device = input_ids.device
+
+        # Normalize shapes: allow [L], [1,L], [B,L], [B,1,L]
+        if input_ids.dim() == 1:
+            input_ids = input_ids.unsqueeze(0)  # [L] -> [1,L]
+        if input_ids.dim() == 3 and input_ids.size(1) == 1:
+            input_ids = input_ids.squeeze(1)    # [B,1,L] -> [B,L]
+        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
+
+                # Top-k
+                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)
+
+                # Top-p
+                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")
+
+                probs = F.softmax(logits, dim=-1)
+
+                if do_sample:
+                    dist = Categorical(probs)
+                    next_token = dist.sample()
+                    log_p = dist.log_prob(next_token)
+                else:
+                    next_token = torch.argmax(probs, dim=-1)
+                    log_p = torch.log(torch.gather(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)      # [B, T]
+        generated_logprobs = torch.cat(generated_logprobs, dim=1)  # [B, T]
+        generated_probs = torch.cat(generated_probs, dim=1) if return_probs else None
+
+        # Use passed tokenizer, fallback to self.tokenizer
+        tok = tokenizer if tokenizer is not None else getattr(self, "tokenizer", None)
+        if tok is None:
+            raise ValueError("Tokenizer must be provided to decode generated tokens.")
+
+        decoded_list = [
+            tok.decode(tok_ids, skip_special_tokens=True)
+            for tok_ids in generated_tokens
+        ]
+        return decoded_list, generated_logprobs, generated_tokens, generated_probs
+
+
+    def ppo_step(
+        self,
+        input_ids: torch.LongTensor,
+        old_log_probs: torch.Tensor,
+        old_action_probs: torch.Tensor,
+        tokenizer,
+        max_new_tokens: int = 50,
+        temperature: float = 1.0,
+        top_k: Optional[int] = 50,
+        top_p: Optional[float] = 0.95,
+        validity_weight: float = 1.0,
+        lipinski_weight: float = 1.0,
+        entropy_weight: float = 0.01,
+        clip_epsilon: float = 0.2,
+        baseline: Optional[torch.Tensor] = None,
+        reward_mode: str = "chemq3",   # "chemq3", "sa", or "mix"
+        reward_mix: float = 0.5,       # used if mixing chemq3 + sa (0..1 weight for chemq3)
+    ) -> Dict[str, Any]:
+
+        # =========================
+        #  PPO-KL BODY  (drop-in)
+        # =========================
+        self.train()
+        self.set_mtp_training(False)
+        if not hasattr(self, 'tokenizer'):
+            self.tokenizer = tokenizer
+
+        # Ensure entropy controller exists
+        if not hasattr(self, 'entropy_controller'):
+            # if you want different defaults, set them when constructing model instead
+            self.entropy_controller = EnhancedEntropyController(
+                min_entropy=0.5,
+                max_entropy=3.0,
+                target_entropy=1.5,
+                adaptation_rate=0.01
+            )
+
+        # --- roll-out ---
+        selfies_list, new_log_probs, token_ids, new_action_probs = self.generate_with_logprobs(
+            input_ids=input_ids,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            do_sample=True,
+            return_probs=True,
+            tokenizer=getattr(self, "tokenizer", None),
+        )
+
+        batch_size = len(selfies_list)
+        device = new_log_probs.device
+
+        # --- rewards: compute depending on mode ---
+        validity_vals: List[float] = []
+        lipinski_vals: List[float] = []
+        total_rewards: List[float] = []
+        sa_rewards: List[float] = []
+
+        for s in selfies_list:
+            if reward_mode == "chemq3":
+                r = compute_comprehensive_reward(s)
+                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(s)
+                sa_rewards.append(sa)
+
+            elif reward_mode == "mix":
+                r = compute_comprehensive_reward(s)
+                sa = compute_sa_reward(s)
+                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 lists -> tensors, handle empty lists safely
+        if reward_mode in ("chemq3", "mix"):
+            rewards = torch.tensor(total_rewards, dtype=torch.float32, device=device)
+        elif reward_mode == "sa":
+            rewards = torch.tensor(sa_rewards, dtype=torch.float32, device=device)
+        else:
+            rewards = torch.zeros(batch_size, dtype=torch.float32, device=device)
+
+        if len(validity_vals) > 0:
+            validity_rewards = torch.tensor(validity_vals, dtype=torch.float32, device=device)
+        else:
+            validity_rewards = torch.zeros(batch_size, dtype=torch.float32, device=device)
+
+        if len(lipinski_vals) > 0:
+            lipinski_rewards = torch.tensor(lipinski_vals, dtype=torch.float32, device=device)
+        else:
+            lipinski_rewards = torch.zeros(batch_size, dtype=torch.float32, device=device)
+
+        # baseline subtraction (broadcast if needed)
+        if baseline is not None:
+            # baseline can be scalar tensor or per-batch; support both
+            if baseline.numel() == 1:
+                rewards = rewards - baseline.to(device)
+            else:
+                rewards = rewards - baseline.to(device)
+
+        # --- probability ratio ---
+        # old_action_probs/new_action_probs expected shape: [B, T, V]
+        # token_ids expected shape: [B, T]
+        old_probs = torch.gather(old_action_probs, 2, token_ids.unsqueeze(2)).squeeze(2).clamp_min(1e-8)
+        new_probs = torch.gather(new_action_probs, 2, token_ids.unsqueeze(2)).squeeze(2).clamp_min(1e-8)
+        log_ratio = new_log_probs - old_log_probs  # shape [B, T]
+        # total_ratio: product of per-step ratios -> exp(sum(log ratio))
+        total_ratio = torch.exp(log_ratio.sum(dim=1))  # shape [B]
+
+        # --- adaptive KL controller (singleton) ---
+        if not hasattr(self, 'kl_controller'):
+            self.kl_controller = AdaptiveKLController()
+        # KL per example: sum over time of old * (log old - log new), averaged over V already via gather
+        # Here compute KL between full distributions if available
+        kl = (old_probs * (torch.log(old_probs) - torch.log(new_probs))).sum(dim=1)  # shape [B]
+        beta = self.kl_controller.update(kl.mean().item())
+
+        # --- PPO-KL loss ---
+        surr1 = total_ratio * rewards
+        surr2 = torch.clamp(total_ratio, 1 - clip_epsilon, 1 + clip_epsilon) * rewards
+        ppo_loss = -torch.min(surr1, surr2).mean()
+        kl_penalty = beta * kl.mean()
+        total_policy_loss = ppo_loss + kl_penalty
+
+        # --- entropy bonus (adaptive) ---
+        # compute token-level entropy averaged across batch/time
+        with torch.no_grad():
+            _probs = new_action_probs.clamp_min(1e-12)
+            per_step_entropy = -(_probs * torch.log(_probs)).sum(dim=-1)  # [B, T]
+            entropy = per_step_entropy.mean()  # scalar tensor
+
+        adaptive_entropy_weight = self.entropy_controller.update_entropy_weight(entropy.item())
+        entropy_bonus = adaptive_entropy_weight * entropy
+        total_loss = total_policy_loss - entropy_bonus
+
+        # regularization (optional)
+        reg_loss = 1e-7 * sum(p.pow(2).sum() for p in self.parameters())
+        total_loss = total_loss + reg_loss
+
+        # prepare return (detach tensors where relevant)
+        avg_sa = None
+        if len(sa_rewards) > 0:
+            avg_sa = float(torch.tensor(sa_rewards, dtype=torch.float32, device=device).mean().item())
+
+        return {
+            'loss': total_loss,
+            'ppo_loss': ppo_loss.item(),
+            'kl_penalty': kl_penalty.item(),
+            'kl_coef': beta,
+            'entropy': float(entropy.item()),
+            'entropy_weight': float(adaptive_entropy_weight),
+            'validity_rate': float(validity_rewards.mean().item()),
+            'lipinski_score': float(lipinski_rewards.mean().item()),
+            'avg_reward': float(rewards.mean().item()),
+            'avg_sa_reward': avg_sa,
+            'generated_selfies': selfies_list,
+            'generated_smiles': [selfies_to_smiles(s) for s in selfies_list],
+            'new_log_probs': new_log_probs.detach(),
+            'new_action_probs': new_action_probs.detach(),
+        }
+
+
+
+
+# ========================
+# CURRICULUM LEARNING MANAGER
+# ========================
+
+class CurriculumManager:
+    def __init__(self, start_len=10, max_len=30, step_increase=5, steps_per_level=30):
+        """
+        Cyclic curriculum:
+        - Gradually increases max_new_tokens from start_len → max_len
+        - After reaching max_len, resets back to start_len and repeats
+        """
+        self.start_len = start_len
+        self.max_len = max_len
+        self.step_increase = step_increase
+        self.steps_per_level = steps_per_level
+        self.step_counter = 0
+        self.current_max_len = start_len
+
+    def get_max_new_tokens(self):
+        return self.current_max_len
+
+    def step(self):
+        self.step_counter += 1
+        if self.step_counter % self.steps_per_level == 0:
+            if self.current_max_len < self.max_len:
+                self.current_max_len = min(self.current_max_len + self.step_increase, self.max_len)
+            else:
+                # Reset cycle
+                self.current_max_len = self.start_len
+                print(f" 🔄 Cycle reset: max_new_tokens -> {self.current_max_len}")
+            if self.current_max_len < self.max_len:
+                print(f" 📈 Curriculum Update: max_new_tokens = {self.current_max_len}")
+        return self.current_max_len
+
+class AdaptiveKLController:
+    """
+    Increases or decreases β so that E[KL] stays ≈ target_kl.
+    """
+    def __init__(self, init_kl_coef: float = 0.1, target_kl: float = 0.01,
+                 kl_horizon: int = 1000, increase_rate: float = 1.5, decrease_rate: float = 0.8):
+        self.kl_coef = init_kl_coef
+        self.target_kl = target_kl
+        self.kl_horizon  = kl_horizon
+        self.inc = increase_rate
+        self.dec = decrease_rate
+        self.buffer = []
+
+    def update(self, kl: float):
+        self.buffer.append(kl)
+        if len(self.buffer) >= self.kl_horizon:
+            avg_kl = sum(self.buffer) / len(self.buffer)
+            self.buffer.clear()
+            if avg_kl > self.target_kl * 1.5:
+                self.kl_coef *= self.inc
+            elif avg_kl < self.target_kl * 0.5:
+                self.kl_coef *= self.dec
+        return self.kl_coef
+    
+
+class EnhancedEntropyController:
+    """
+    More sophisticated entropy control with dynamic targets and temperature scheduling.
+    """
+    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  # Starting weight
+        
+    def update_entropy_weight(self, current_entropy: float) -> float:
+        """
+        Dynamically adjust entropy weight based on current entropy levels.
+        """
+        self.entropy_history.append(current_entropy)
+        
+        # Keep rolling window
+        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 entropy too low, increase weight to encourage exploration
+            if avg_entropy < self.target_entropy * 0.8:
+                self.entropy_weight = min(0.05, self.entropy_weight * 1.1)
+            # If entropy too high, decrease weight
+            elif avg_entropy > self.target_entropy * 1.2:
+                self.entropy_weight = max(0.001, self.entropy_weight * 0.95)
+                
+        return self.entropy_weight
+    
+    def compute_entropy_reward(self, entropy: float) -> float:
+        """
+        Reward function for entropy - prefer target range.
+        """
+        if self.min_entropy <= entropy <= self.max_entropy:
+            # Gaussian reward centered at target
+            distance = abs(entropy - self.target_entropy)
+            max_distance = max(self.target_entropy - self.min_entropy, 
+                             self.max_entropy - self.target_entropy)
+            return np.exp(-(distance / max_distance) ** 2)
+        else:
+            return 0.1  # Small penalty for being outside range

FastChemTokenizerHF.py

@@ -0,0 +1,769 @@
+import torch
+import json
+import os
+from typing import List, Union, Optional, Tuple, Dict, Any
+from transformers.tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase
+from transformers.utils import PaddingStrategy, TensorType
+from functools import lru_cache
+
+
+class TrieNode:
+    __slots__ = ['children', 'token_id']
+    def __init__(self):
+        self.children = {}
+        self.token_id = None  # If set, this node completes a valid token
+
+
+class FastChemTokenizer(PreTrainedTokenizerBase):
+    """
+    Fully HuggingFace API compatible tokenizer for chemical representations.
+    """
+
+    vocab_files_names = {"vocab_file": "vocab.json"}
+
+    def __init__(
+        self,
+        token_to_id=None,
+        vocab_file=None,
+        model_max_length=512,
+        padding_side="right",
+        truncation_side="right",
+        chat_template=None,
+        **kwargs
+    ):
+        # Handle vocab loading
+        if token_to_id is None and vocab_file is None:
+            raise ValueError("Either token_to_id or vocab_file must be provided")
+
+        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()}
+
+        # Precompute max token length for possible use & clarity
+        self.max_token_len = max(len(t) for t in token_to_id.keys()) if token_to_id else 0
+
+        # Build trie for fast longest-match lookup
+        self.trie_root = self._build_trie(token_to_id)
+
+        # Validate required special tokens
+        required_special_tokens = ["<s>", "</s>", "<pad>", "<unk>", "<mask>"]
+        for tok in required_special_tokens:
+            if tok not in token_to_id:
+                raise KeyError(f"Required special token '{tok}' not found in vocab.")
+
+        # ✅ Assign special token IDs explicitly
+        self.bos_token_id = token_to_id["<s>"]
+        self.eos_token_id = token_to_id["</s>"]
+        self.pad_token_id = token_to_id["<pad>"]
+        self.unk_token_id = token_to_id["<unk>"]
+        self.mask_token_id = token_to_id["<mask>"]
+
+        # Special tokens
+        bos_token = "<s>"
+        eos_token = "</s>"
+        pad_token = "<pad>"
+        unk_token = "<unk>"
+        mask_token = "<mask>"
+
+        # Initialize parent class with all required parameters
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            sep_token=None,
+            pad_token=pad_token,
+            cls_token=None,
+            mask_token=mask_token,
+            additional_special_tokens=[],
+            model_max_length=model_max_length,
+            padding_side=padding_side,
+            truncation_side=truncation_side,
+            chat_template=chat_template,
+            **kwargs,
+        )
+
+    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]:
+        """Core encoding logic using Trie — no caching."""
+        tokens = text
+        result_ids = []
+        i = 0
+        n = len(tokens)
+
+        while i < n:
+            node = self.trie_root
+            j = i
+            last_match_id = None
+            last_match_end = 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 = node.token_id
+                    last_match_end = j
+
+            if last_match_id is not None:
+                result_ids.append(last_match_id)
+                i = last_match_end
+            else:
+                tok = tokens[i]
+                result_ids.append(self.token_to_id.get(tok, self.unk_token_id))
+                i += 1
+
+        return result_ids
+
+    def _tokenize(self, text: str, **kwargs) -> List[str]:
+        token_ids = self._encode_core(text.strip())
+        return [self.id_to_token[tid] for tid in token_ids]
+
+    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)
+
+    # ✅ Public methods
+    def convert_tokens_to_ids(self, tokens: Union[str, List[str]]) -> Union[int, List[int]]:
+        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]]) -> Union[str, List[str]]:
+        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:
+        """SMILES-style decoding: no spaces between tokens."""
+        return "".join(tokens)
+
+    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 decode(
+        self,
+        token_ids: Union[List[int], torch.Tensor],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = None,
+        **kwargs
+    ) -> str:
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+
+        if skip_special_tokens:
+            special_ids = {
+                self.bos_token_id,
+                self.eos_token_id,
+                self.pad_token_id,
+                self.mask_token_id,
+            }
+        else:
+            special_ids = set()
+
+        tokens = []
+        for tid in token_ids:
+            if tid in special_ids:
+                continue
+            token = self.id_to_token.get(tid, self.unk_token)
+            tokens.append(token)
+
+        return "".join(tokens)
+    
+    def batch_decode(
+        self,
+        sequences: Union[List[List[int]], torch.Tensor],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = None,
+        **kwargs
+    ) -> List[str]:
+        """Batch decode sequences."""
+        if isinstance(sequences, torch.Tensor):
+            sequences = sequences.tolist()
+        
+        return [
+            self.decode(
+                seq, 
+                skip_special_tokens=skip_special_tokens,
+                clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+                **kwargs
+            ) 
+            for seq in sequences
+        ]
+
+    def decode_with_trace(self, token_ids: List[int]) -> None:
+        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)
+            print(f"  [{i:03d}] ID={tid:5d} → '{token}'")
+
+    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, PaddingStrategy] = 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, TensorType]] = 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:
+        """
+        Main callable method that handles both single and batch inputs.
+        """
+        # Handle defaults
+        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, PaddingStrategy] = 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, TensorType]] = 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()))
+
+        if text_pair is not None:
+            ids_b = list(self._cached_encode_str(text_pair.strip()))
+        else:
+            ids_b = None
+
+        input_ids = []
+        token_type_ids = []
+
+        if add_special_tokens:
+            input_ids.append(self.bos_token_id)
+            token_type_ids.append(0)
+            if ids_b is not None:
+                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)
+
+                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.extend(ids_a)
+                token_type_ids.extend([0] * len(ids_a))
+                input_ids.append(self.eos_token_id)
+                token_type_ids.append(0)
+        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))
+
+        # Handle truncation
+        if truncation and len(input_ids) > max_length:
+            input_ids = input_ids[:max_length]
+            token_type_ids = token_type_ids[:max_length]
+
+        # Handle padding
+        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 = [1 if tid != self.pad_token_id else 0 for tid in input_ids]
+
+        encoded_dict = {
+            "input_ids": input_ids,
+        }
+        
+        if return_attention_mask:
+            encoded_dict["attention_mask"] = attention_mask
+        
+        if return_token_type_ids:
+            encoded_dict["token_type_ids"] = token_type_ids
+            
+        if return_special_tokens_mask:
+            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
+            ]
+            encoded_dict["special_tokens_mask"] = special_tokens_mask
+            
+        if return_length:
+            encoded_dict["length"] = len([tid for tid in input_ids if tid != self.pad_token_id])
+
+        if return_tensors == "pt":
+            output = {}
+            for k, v in encoded_dict.items():
+                tensor = torch.tensor(v, dtype=torch.long)
+                if tensor.ndim == 1:
+                    tensor = tensor.unsqueeze(0)
+                output[k] = tensor
+        else:
+            output = encoded_dict
+
+        return BatchEncoding(output, tensor_type=return_tensors)
+
+    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, PaddingStrategy] = 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, TensorType]] = 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:
+        all_input_ids = []
+        all_attention_masks = []
+        all_token_type_ids = []
+        all_special_tokens_masks = []
+        all_lengths = []
+
+        for item in batch_text_or_text_pairs:
+            if isinstance(item, tuple):
+                text, text_pair = item
+            else:
+                text, text_pair = item, None
+
+            encoded = self.encode_plus(
+                text=text,
+                text_pair=text_pair,
+                add_special_tokens=add_special_tokens,
+                padding=False,  # We'll handle batch padding later
+                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=None,  # Don't convert to tensors yet
+                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
+            )
+            
+            all_input_ids.append(encoded["input_ids"])
+            if "attention_mask" in encoded:
+                all_attention_masks.append(encoded["attention_mask"])
+            if "token_type_ids" in encoded:
+                all_token_type_ids.append(encoded["token_type_ids"])
+            if "special_tokens_mask" in encoded:
+                all_special_tokens_masks.append(encoded["special_tokens_mask"])
+            if "length" in encoded:
+                all_lengths.append(encoded["length"])
+
+        batched = {
+            "input_ids": all_input_ids,
+        }
+        
+        if all_attention_masks:
+            batched["attention_mask"] = all_attention_masks
+        if all_token_type_ids:
+            batched["token_type_ids"] = all_token_type_ids
+        if all_special_tokens_masks:
+            batched["special_tokens_mask"] = all_special_tokens_masks
+        if all_lengths:
+            batched["length"] = all_lengths
+
+        # Handle batch padding
+        if padding == True or padding == "longest":
+            max_len = max(len(ids) for ids in all_input_ids)
+            for key in batched:
+                if key in ["input_ids", "attention_mask", "token_type_ids", "special_tokens_mask"]:
+                    padded_seqs = []
+                    for seq in batched[key]:
+                        pad_len = max_len - len(seq)
+                        if pad_len > 0:
+                            if key == "input_ids":
+                                padding_value = self.pad_token_id
+                            else:
+                                padding_value = 0
+                            
+                            if self.padding_side == "right":
+                                padded_seq = seq + [padding_value] * pad_len
+                            else:
+                                padded_seq = [padding_value] * pad_len + seq
+                        else:
+                            padded_seq = seq
+                        padded_seqs.append(padded_seq)
+                    batched[key] = padded_seqs
+
+        if return_tensors == "pt":
+            def to_tensor_list(lst):
+                return [torch.tensor(item, dtype=torch.long) for item in lst]
+            
+            for key in ["input_ids", "attention_mask", "token_type_ids", "special_tokens_mask"]:
+                if key in batched:
+                    batched[key] = torch.nn.utils.rnn.pad_sequence(
+                        to_tensor_list(batched[key]),
+                        batch_first=True,
+                        padding_value=self.pad_token_id if key == "input_ids" else 0
+                    )
+            
+            # Handle non-sequence data
+            if "length" in batched:
+                batched["length"] = torch.tensor(batched["length"], dtype=torch.long)
+
+        return BatchEncoding(batched, tensor_type=return_tensors)
+
+    def pad(
+        self,
+        encoded_inputs,
+        padding: Union[bool, str, PaddingStrategy] = True,
+        max_length: Optional[int] = None,
+        pad_to_multiple_of: Optional[int] = None,
+        return_attention_mask: Optional[bool] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        verbose: bool = True,
+    ) -> BatchEncoding:
+        """Pad encoded inputs."""
+        # This is a simplified version - full implementation would be more complex
+        return encoded_inputs
+
+    # Save/Load methods
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """Save vocabulary to files."""
+        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],
+        legacy_format: bool = True,
+        filename_prefix: Optional[str] = None,
+        push_to_hub: bool = False,
+        **kwargs
+    ):
+        """Save tokenizer to directory."""
+        if not os.path.exists(save_directory):
+            os.makedirs(save_directory)
+
+        # Save vocabulary
+        vocab_files = self.save_vocabulary(save_directory, filename_prefix)
+        
+        # Save tokenizer config
+        tokenizer_config = {
+            "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,
+            }
+        }
+        
+        config_file = os.path.join(save_directory, "tokenizer_config.json")
+        with open(config_file, "w", encoding="utf-8") as f:
+            json.dump(tokenizer_config, f, ensure_ascii=False, indent=2)
+
+        print(f"✅ Tokenizer saved to: {save_directory}")
+        
+        return (save_directory,)
+
+    @classmethod
+    def from_pretrained(
+        cls, 
+        pretrained_model_name_or_path: Union[str, os.PathLike],
+        *init_inputs,
+        **kwargs
+    ):
+        """Load tokenizer from pretrained directory or hub."""
+        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")
+            
+            # Load config if available
+            config = {}
+            if os.path.exists(config_file):
+                with open(config_file, "r", encoding="utf-8") as f:
+                    config = json.load(f)
+            
+            # Merge config with kwargs
+            merged_config = {**config, **kwargs}
+            
+            return cls(vocab_file=vocab_file, **merged_config)
+        else:
+            raise NotImplementedError("Loading from HuggingFace Hub not implemented yet")
+
+    def get_special_tokens_mask(
+        self, 
+        token_ids_0: List[int], 
+        token_ids_1: Optional[List[int]] = None, 
+        already_has_special_tokens: bool = False
+    ) -> List[int]:
+        """Get special tokens mask."""
+        if already_has_special_tokens:
+            return [
+                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 token_ids_0
+            ]
+        
+        mask = [1]  # BOS
+        mask.extend([0] * len(token_ids_0))  # Token sequence
+        mask.append(1)  # EOS
+        
+        if token_ids_1 is not None:
+            mask.extend([0] * len(token_ids_1))  # Second sequence
+            mask.append(1)  # EOS
+            
+        return mask
+
+    def create_token_type_ids_from_sequences(
+        self, 
+        token_ids_0: List[int], 
+        token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """Create token type IDs for sequences."""
+        sep = [self.eos_token_id]
+        cls = [self.bos_token_id]
+        
+        if token_ids_1 is None:
+            return len(cls + token_ids_0 + sep) * [0]
+        
+        return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
+
+    def build_inputs_with_special_tokens(
+        self, 
+        token_ids_0: List[int], 
+        token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """Build inputs with special tokens."""
+        if token_ids_1 is None:
+            return [self.bos_token_id] + token_ids_0 + [self.eos_token_id]
+        
+        return ([self.bos_token_id] + token_ids_0 + [self.eos_token_id] + 
+                token_ids_1 + [self.eos_token_id])
+
+
+class FastChemTokenizerSelfies(FastChemTokenizer):
+    """
+    SELFIES variant that handles whitespace-separated tokens.
+    Uses trie-based longest-match encoding (same as original working version).
+    """
+
+    def _encode_core(self, text: str) -> List[int]:
+        """Trie-based encoding for SELFIES with fragment + atom vocab."""
+        result_ids = []
+        i = 0
+        n = len(text)
+
+        while i < n:
+            if text[i].isspace():  # skip literal whitespace
+                i += 1
+                continue
+
+            node = self.trie_root
+            j = i
+            last_match_id = None
+            last_match_end = i
+
+            # Traverse trie character by character (including spaces if part of vocab key)
+            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 = node.token_id
+                    last_match_end = j
+
+            if last_match_id is not None:
+                result_ids.append(last_match_id)
+                i = last_match_end
+            else:
+                # Fallback: encode one char as unk or atom
+                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:
+        """SELFIES decoding: join tokens with spaces (preserve original format)."""
+        return " ".join(tokens)
+
+    def decode(
+        self,
+        token_ids: Union[List[int], torch.Tensor],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = None,
+        **kwargs
+    ) -> str:
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+
+        if skip_special_tokens:
+            special_ids = {
+                self.bos_token_id,
+                self.eos_token_id,
+                self.pad_token_id,
+                self.mask_token_id,
+            }
+        else:
+            special_ids = set()
+
+        tokens = []
+        for tid in token_ids:
+            if tid in special_ids:
+                continue
+            token = self.id_to_token.get(tid, self.unk_token)
+            tokens.append(token)
+
+        return " ".join(tokens)   # ✅ preserve spaces

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 gbyuvd
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

config.json

@@ -0,0 +1,34 @@
+{
+  "training": {
+    "batch_size": 16,
+    "num_epochs": 1,
+    "learning_rate": 5e-5,
+    "weight_decay": 0.01,
+    "gradient_accumulation_steps": 4,
+    "tokenize_batch_size": 100,
+    "train_split_ratio": 0.8,
+    "val_split_ratio": 0.1,
+    "test_split_ratio": 0.1,
+    "include_for_metrics": ["input_ids", "attention_mask", "labels"]
+  },
+  "model": {
+    "max_position_embeddings": 512,
+    "hidden_size": 320,
+    "num_hidden_layers": 6,
+    "num_attention_heads": 4,
+    "num_key_value_heads": 2,
+    "head_dim": 64,
+    "intermediate_size": 1280,
+    "sliding_window": 16,
+    "rope_theta": 10000.0,
+    "attention_dropout": 0.1
+  },
+  "generation": {
+    "max_length": 64,
+    "top_k": 50,
+    "top_p": 0.9,
+    "temperature": 1,
+    "do_sample": true,
+    "num_return_sequences": 3
+  }
+}

demo_test_mtpresult.ipynb

@@ -0,0 +1,190 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "4ff9650b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[  0, 379,   1]])\n",
+      "tensor([[1, 1, 1]])\n",
+      "cuda:0\n"
+     ]
+    }
+   ],
+   "source": [
+    "from FastChemTokenizerHF import FastChemTokenizerSelfies\n",
+    "# --- Load the tokenizer ---\n",
+    "tokenizer = FastChemTokenizerSelfies.from_pretrained(\"./selftok_core\")\n",
+    "\n",
+    "# Test it\n",
+    "out = tokenizer(\"[C]\", return_tensors=\"pt\")\n",
+    "print(out.input_ids)          # ← Attribute access works\n",
+    "print(out.attention_mask)     # ← Also works\n",
+    "out = out.to(\"cuda\")          # ← Moves all tensors to GPU\n",
+    "print(out.input_ids.device)   # ← Should be cuda:0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "d16aeaf7",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Model has 9,854,851 trainable parameters.\n",
+      "Input shape: torch.Size([2, 32])\n",
+      "Logits shape: torch.Size([2, 32, 782])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "from ChemQ3MTP import ChemQ3MTP\n",
+    "# --- Initialize model from scratch ---\n",
+    "\n",
+    "model = ChemQ3MTP.from_pretrained('./enhanced-qwen3-final')\n",
+    "\n",
+    "# --- Print model parameter count ---\n",
+    "def count_parameters(model):\n",
+    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
+    "\n",
+    "print(f\"Model has {count_parameters(model):,} trainable parameters.\")\n",
+    "\n",
+    "# --- Quick forward pass sanity check ---\n",
+    "batch_size, seq_len = 2, 32\n",
+    "dummy_input = torch.randint(\n",
+    "    low=0,\n",
+    "    high=len(tokenizer),\n",
+    "    size=(batch_size, seq_len),\n",
+    "    dtype=torch.long,\n",
+    ")\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    outputs = model(dummy_input)\n",
+    "    logits = outputs.logits\n",
+    "\n",
+    "print(f\"Input shape: {dummy_input.shape}\")\n",
+    "print(f\"Logits shape: {logits.shape}\")  # should be [batch_size, seq_len, vocab_size]\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "105b47a0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Branch2] [=Branch1] [Branch1] [C] [=Branch1] [C] [=O] [N] [C] [C] [N] [C] [=Branch1] [C] [=O] [C] [N] [C] [=Branch1] [C] [=O] [NH1] [C] [=Ring2] [Ring1] [=Branch1] [=C] [Branch2] [Ring1] [C] [C] [C] [O] [S] [=Branch1] [C] [=O] [=Branch1] [C] [=O] [C] [=C] [C] [=C] [C] [Branch1] [=Branch2] [N] [C] [C] [N] [Branch1] [C] [C] [C] [=Branch1] [C] [=O] [O] [C] [=C] [Ring1] [=C] [Ring1] [#Branch1] [C] [Branch2] [Ring1] [O] [C] [C] [O] [C] [=N]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Generate SELFIES\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "model.to(device)\n",
+    "input_ids = tokenizer(\"<s>\", return_tensors=\"pt\").input_ids.to(device)\n",
+    "gen = model.generate(input_ids, max_length=256, top_k=50, temperature=1, do_sample=True, pad_token_id=tokenizer.pad_token_id)\n",
+    "print(tokenizer.decode(gen[0], skip_special_tokens=True))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "b041d311",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "C1(=O)NCCNC(=O)CNC(=O)[NH1]C1C(CCOS(=O)(=O)C=C2C=CCNCCN(C)C)(O)OC=C2\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Manually convert it to SMILES\n",
+    "import selfies as sf\n",
+    "\n",
+    "test = tokenizer.decode(gen[0], skip_special_tokens=True)\n",
+    "test = test.replace(' ', '')\n",
+    "print(sf.decoder(test))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "f1608fa0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "C=1=NC2=CC=CC=C2N=1\n"
+     ]
+    },
+    {
+     "data": {
+      "image/jpeg": "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/2wBDAQkJCQwLDBgNDRgyIRwhMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjL/wAARCAEsASwDASIAAhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwD3+iiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigBskiRRtJI6oijLMxwAPc02G4huY/MgmjlTON0bBh+YrnfiJ/wAk71//AK83/lXk/wAK9ZuvB2oabYak/wDxJ/EMfm2sh4WOcEoR+OAD9UNerhcseJwk68Je9F6Lvpd287a28iJTtKx71PcQWsfmXE0cSZxukYKM/U08EMAQQQeQRXmfx2/5J6n/AF/RfyavQtL/AOQRZf8AXBP/AEEVzVMKoYWGIv8AE2relv8AMal7zRbooorjKCiiigAooooAKKKKACiiigAooooAKKKiurq3sraS5up44IIl3PLKwVVHqSeBQBLWVr/iXR/C+nm91m/itYei7zlnPoqjlj7CuOn8e6v4qnksfh/pwuI1YpLrV6pS1jPfYOsh+n5EVpeH/hxYadqA1nW7qbXteOCb29GRGfSJOiD07jtigDBk8T/EnxG7aj4X8PWthpUI3xJqx2zXw9AuRsBGCM4H+0elb3hv4jafq94NI1a2l0PXxw2n3vylz6xscBwe2OfbvXaVj+IfC+jeKrD7HrNhFcxjlGPDxn1Vhyp+lAGxRXmv2Xxr8P8AmyebxX4fT/l3lYC+t1/2W/5agenXsMV1nhnxlofi23aTSrwNLHxNbSDZNCe4ZDyOeM9PegDeooooAKKKKACiiigAooooAKKKKACiiigAooooA5n4if8AJO9f/wCvN/5VxWneEU8YfArSbNAFvoYWmtJOmJA7cZ9D0/I9q7rx3bT3ngTWra1hknnktHVI41LMxx0AHU1X+G9nc2Hw90e1vLeW3uI4mDxSoVZTvbqDyK9nD4mVDAKdN2kqia/8Bf4GbV52fY8k8T+LZPE3waWC/JXV9Ov4re8R+GJCuA5Hvg59wa9jh8T6Hpi6dpl/qlra3klpHIkcz7NykYByeOoPGc15X8YPh/qLav8A2zoFlc3Md+QLy3toy5Eg5D7R2Pr6j/ar146Hp2raJaWuradb3KrAgKXEQYqdo9ehrtzGeDlhqMo/DJydla8W7aeid7baEwUuZmsrK6hkYMpGQQcgilrhW+GkOnM0vhbXNS0KTORFHKZrcn3jfr+dJ/aXxB0Ef6fpNj4gtl/5bWEnkzY9SjcE+y15H1OnU/gVE/KXuv8AH3f/ACY05mt0d3RXHaf8TfDd1cC1vp59Hve9tqkRgYfifl/WuuiljniWWGRJI2GVdDkEexFc1bDVqDtVi16jTT2H0UUViMKKKKACiiigAorm/E/jnRPCgSK9naa/l4gsLVfMuJiegCD19TgVzP8AY3jHx98/iCd/Dmgv00uzfNzOv/TWT+EH+6PoR3oA0tb+JFpb6i2i+G7OXxBrg4NvaH91D2zLL91QP/14qnbfD7UfElzHqPxA1EagynfFpFqSlnAfcdZD7n6ciuy0Pw/pPhvT1sNHsIbO3X+GMcsfVj1Y+5JNVPEvjDRPCVqs2rXqxu/EVug3zTH0RByee/T3oA2YIIbaBIIIkihjUKkcahVUDsAOgrkvEfxF0zRr7+yNPgm1rXm4TTrEbmU/9NG6IPXPPtWN5Xjb4gcztN4T8Pv/AMskP+n3C+56RA/n9RXY+HPCui+FLH7Jo1jHbq3MknWSU+rMeSaAMXw3pXjC51hdc8T6sluAjLFo1jjyYwe8jHl2Ht0PQ4OK7OiigArlPE3w/wBH8R3K6ghm03WY+YtTsW8uZT/tEfeHse3cV1dFAHmy+LPE3gdhB41s/wC0dKBwuu6fETtHrPEOV+o4+td9pmqWGs2Ed9pt3Dd2sgyssLhgfb2PtVplDKVYAgjBB71wWp/Dc2V/JrHgnUDoGpv80kCjdaXJ9Hj6D6r09M0Ad9RXAad8R20++j0jxxp50LUXO2O5J3WdwfVJP4fo3T17V3ysrqGVgykZBByCKAFooooAKKKKACiiigAooooAKa7rGjO7BUUEszHAA9TTq4DxZeXHizXR4I0qVkgCiTWbqP8A5ZRHpED/AHn/AJevNdGGw7rzteyWrfZd/wDLu9BSdkWfDGran4s8SXWuQ3EkHhu3VrWzixgXjZ+aU57AjA//AFg9tUFnZ2+n2UNnaRLDbwII441HCqBgCuI1K+u/GPjFNC0q5mg0nSZVm1O7hcqZJQcrArD3GW/+tzu4rFVW4LlhFfcl37tv72+wvhXmd9RRRXAUFFFFAFTUNLsNWtzb6jZW93D/AHJ4w4/WuQufhhZW2+Twzq+peHp2bfi0mLxE8/ejbII56ZFd1RWsa9WMHTUnyvdX0+4Vle5559s+Jvh0f6TYaZ4otV/5aWr/AGW5I9Sp+Q/Ras2PxZ8OSXC2esC80C+P/LDVoDD+T/dx7kiu6qtfafZanbNbX9pBdQN1injDqfwPFZDJLa6t7y3S4tZ4p4XGVkicMrD2I4NS1wNz8JtGgme68OXuo+HLtjktp1wwjY/7UZJBHsMVi65q/wASPBcdhbTX+iazHqV7Hp9rdzQNDKksmdpdE+UrwenNAHpGs67pfh7T3v8AV76Gztk6vK2Mn0A6k+w5rhf7e8XePf3fhm2bQNDfg6vex5nmX1hi7D0Y/hg1o6N8N7ddQTWvFV7J4h1ocrJcj9xB7RRfdA9/x4ruCQqlmICgZJPagDmvDHgTRPCrPcW0UlzqUvM+o3bebcSk9csen0GK3NR1Kx0ixkvdRu4bW1jGXlmcKo/E/wAq4vU/iR9tvpNH8E6edf1NflkmQ4tLc+ry9D9B19c0mnfDh9Rvo9X8c6gdd1BDujtcbbO2Pokf8X1br3HegCs3i/xJ43c2/giy+xaYTtfXtQjIUj1hiPLn3PHqBW34a+Huk+H7ptSnebVdbk5l1O+bzJSf9nPCD2HbjJrrFVUUKqhVUYAAwAKWgAooooAKKKKACiiigAooooAq6jpljq9jJZajaQ3VrIMPFMgZT+B7+9cC3hDxJ4Ic3Hgi+N7pgOW0HUJCVA9IZTyh9jx6k16RRQByXhr4haT4hum02dZtK1uPiXTL5fLlB/2c8OPcducCutrC8S+D9D8W2qw6tZLI6cxXCHZNCfVHHI57dPauP0i68U+DfHej+FtV1WPW9I1cXH2K5nBF1B5Ue8hyOG6gZOSevGMUAem0UUUAFFFFABRRVHWNXstB0i61TUZhDaW0Zkkc+noPUk8AdyaAMfxp4mk0DTorfT4vtGtag/kWFuP4nPVj/sr1P4VP4Q8Mx+GNG+ztKbi+ncz3t03LTTNyxz6dh/8ArrD8Eadd61eyeONct2hvb2PZYWrnP2S27f8AAm6n6+5FdP4h16z8NaJcanek+XEvyxr96R+yqPU//X7V3ymvZxw1DXms2+76L0X4u77E21uzE8b+IbuyW20DQ8Pr+qEpB/07x/xTN6ADOPf1xitjwz4etPC+hwaZaZYJ80srfemkP3nb3J/oO1YvgjQLyJrnxLryg67qmGdD0tYf4YV9MDGff6ZrsaMVONOH1ak7pfE+8v8AJbL5vqEVf3mFFFFcBQUUUUAFFFFABRRRQAV5/wDFT7ngz/sabH/2evQK8/8Aip9zwZ/2NNj/AOz0AegVxviLwTeeLNZYatrlwPDqqu3SrUeV5rfxea4OWXPYY/DHPZUUAU9M0qw0WwjsdMs4bS1jGFihQKB7+59+tXKKKACiiigAooooAKKKKACiiigAooooAKhvJ/stlPcbd3lRs+3OM4GcVNVPV/8AkC33/XvJ/wCgmgCh4R18+KPCmna21sLY3ke/yQ+/ZyRjOBnp6VzXi3/kr3w6/wC4n/6TrVz4S/8AJK/D/wD17n/0Nqp+Lf8Akr3w6/7if/pOtAHoFFFFABRRRQAV5hJ/xdHxj5Q+fwhoU/zn+HULsdvdE/I+4PGh481q+1K/t/A3h6bZqmoJuvLlefsVr0Zz6M3QD37ZBrr9D0Wx8O6La6Tp0IitbZAiL3PqT6knJJ9TQBed0jjZ3ZURRksTgADvXmmjo3xL8Xr4iuFJ8MaPKU0qJhxdzjhpyO6g8L/TkVN4zvrnxfr6+AdHmeOHaJdcvI/+WEB6RA/33/l6jOO/sbG20ywgsbKFYba3QRxRoOFUDAFAFiiiigAooooAKKKKACiiigAooooAK8/+Kn3PBn/Y02P/ALPXoFef/FT7ngz/ALGmx/8AZ6APQKKKKACiiigAooooAKKKKACiiigAooooAKKKKACqer/8gW+/695P/QTVyqer/wDIFvv+veT/ANBNAHL/AAl/5JX4f/69z/6G1U/Fv/JXvh1/3E//AEnWrnwl/wCSV+H/APr3P/obVT8W/wDJXvh1/wBxP/0nWgD0CiiigArD8Ya+/hbwlqOtJaPdvaRb1hTuSQMn0UZyT6A1uU10SSNkdQyMCGVhkEHsaAOK+GWiLa+Hhr91cre6vroW8vLsc53DKxr6KoOMeufoNLx/4jn8KeCtR1e1hEtxEqrHu+6rMwUM3sM5/CuX0Z3+Gvi5fDlwxHhjV5S+lSseLWc8tbk9gTyv9eSPR7u0gvrOa0uolmt5kMckbjIZSMEGgDnvAnhmDw14eRVuRe3t6ftV7fZ3G5lfktnuvPHt9TXT15t4Ru5/A/iT/hBNUld9Pn3S6FdyHO6Pq0DH+8vb2+oFek0AFFFFABRRRQAUUUUAFFFFABRRRQAV5/8AFT7ngz/sabH/ANnr0CvP/ip9zwZ/2NNj/wCz0AegUUUUAFFFFABRRRQAUUUUAFFFFABRRRQAUUUUAFU9X/5At9/17yf+gmrlU9X/AOQLff8AXvJ/6CaAOX+Ev/JK/D//AF7n/wBDaqfi3/kr3w6/7if/AKTrVz4S/wDJK/D/AP17n/0Nqp+Lf+SvfDr/ALif/pOtAHoFFFFABRRRQBkeJ/Dlj4r8P3OkX6nypl+SRfvROPuuvuD/AId65/wD4jvp2u/C3iFgPEOk4WR+13D/AATL65GM+/1wO3rivH3hq9vVtfEfh/CeI9IzJb+lzH/HA3qCM49/TOaANTxn4Vg8XaA9i0ht7uJhPZXS8NbzLyrg/wA/aqXgPxVPr1hcafq0Yt/EOlv5GoW/TLdpF/2WHI/wxWr4V8SWfizw9batZZVZRtkib70Mg4ZG9wf8e9c5430K7stUtfG2gof7T09Cl5Ag/wCP216shHdh1B/ngCqhBzkordgd5RVHRtXs9e0i21OwlEltcJvQ9x6g+hByCPar1EoyhJxkrNAFFFFSAUUUUAFFFFABRRRQAV5/8VPueDP+xpsf/Z69Arz/AOKn3PBn/Y02P/s9AHoFFFFABRRRQAUUUUAFFFFABRRRQAUUUUAFFFFABVPV/wDkC33/AF7yf+gmrlU9X/5At9/17yf+gmgDl/hL/wAkr8P/APXuf/Q2qn4t/wCSvfDr/uJ/+k61c+Ev/JK/D/8A17n/ANDaqfi3/kr3w6/7if8A6TrQB6BRRRQAUUUUAFFFFAHnmuJJ4B8UHxNbKx0HUXWPVoVGRBIeFnA9Ozf1yMegxyJLGkkbq8bgMrKcgg9CDUd1awX1pNa3USywTIY5I2GQykYINcL4Wup/B+vnwVqcrPZyhpdFupD9+PqYSf7y9vb04Fei/wDa6N1/Egvviv1j/wCk+hHwvyIZP+LceLPOHy+Ftam/eD+Gxuj39kb9PbHPo1U9V0u01rSrnTb+IS2twhSRT6eo9COoPqK5HwXql3pGpS+CtclL3lom/T7l/wDl7tu3/Al6Ee3sTRU/2ul7VfHFa+a/m9Vs/Kz7gvdduh3VFFFecWFFFFABRRRQAUUUUAFef/FT7ngz/sabH/2evQK8v+LGvaRHceE7R9Us1uLfxHZ3E0RmXdFGu7c7DPyqMjk0AeoUxZY2kaNXUumCyg8rnpkUsciSxrJG6ujgMrKcgg9CDXIeJfAFvq+p/wBu6TqFzo3iBVCi+tmyJABgLIh4daAOxorzm38e6t4WnjsPiBpwtkYhItas1L2kp7bx1jJ9/wBBXoNtc297bR3NrPHPBINySxOGVh6gjg0AS0UUUAFFFFABRRRQAUUUUAFFFFABVPV/+QLff9e8n/oJq5RQBxfwl/5JX4f/AOvc/wDobVT8W/8AJXvh1/3E/wD0nWu/ACgBQAB2Fec+KL21n+NXgG0iuYpLm3GoGaJXBaMNbjbuHUZwcZ9KAPR6KKKACiiigAooooAKwfF3hqHxRojWhkMF3Ewms7leGgmX7rA/z9q3qK0pVZ0pqpB2aE1dWZy/grxNNrljPZanGINc05/Iv4OnzdnX/ZYcirviDwxaeIHsJ5ZZra7sLhZ7e5gIDpz8y8jow4IqvrVno+iXtx40uIZxc2dmySmBsebHweVyAxGOMn+QwvhHxppPjWwmutLMy+TJskinUK68ZBIBPB5wc9jXbNTu8XhotRW/ZN7r07X6OxK/lkdFRWH4r8VWHg7Rxqeox3EkBlWLFuoZskEjgkccetbEEy3FvFOgIWRA4B64IzXE6U1BVGvdeifpuVdXsSUUdBk1y+r/ABD8L6NJ5M2qRT3WcC2tP30hPphc4P1xTpUKtaXLSi5PyVwbS3Ooorg/+Em8Z67xoPhYadA3S71qTYf+/S/NSSfD3UddiI8WeKtQvUYgtaWJ+ywY/ukLyw574Nb1MG6UW6k4p9r3f4XS+bQlK+xq678RPCnh1jFf6zbm5Bx9mgPnSk+m1ckfjisT/hM/GHiAY8MeDZbaFvu3uuv5C/Xyh85HuDXUaH4Q8PeGkC6Po9paMBjzEjzIR7ucsfxNbdcZR55/wgHiHXlz4t8ZXssTdbHSlFrD/ulh8zj64rb074c+D9KspLS28PWBjkXbIZohK7j3Z8n9a6iigDzeTwVr/g2RrrwHf+bY53PoOoSFoj6+U55Q+xOM9T2rY8OfETS9bvTpN9DNo2upw+nXw2OT/sN0cemOe+K7CsXxH4T0TxZZfZdYsY5wv+rlHyyRH1RxyP8AOaANa4t4bu3kt7mGOaGRdrxyKGVh6EHgivP7n4f6l4auZNR8AakLHcS8uj3ZL2cx77e8ZPqPpwKh2eNvh/8AcM3izw8n8Lf8f9uvsekoH5/QV13hrxfoni20M+kXqSsn+tgb5ZYj6Oh5H8vQ0AYmhfEizutRXRfEVnL4f13oLW7P7ub3ik+6w/yM129Zmu+HtJ8S6c1hrFhDd256LIOVPqp6qfcVxH9keMvAHzaFPJ4l0FOum3b4uoF9IpP4gP7p+gHegD0qiuc8L+ONE8WI6WFw0V7FxPY3K+XPCR1DIf5jIro6ACiiigAooooAKKKjuLiG0t5Li5mjhhjXc8kjBVUepJ4AoAkrM13xFpHhnT2vtYv4bS3HQyHlj6Ko5Y+wFcbc+P8AU/EtzJp/w/04Xu1tkusXYKWcJ77e8hHoP1FX9C+G9naaguteIbyXxBrvX7VdgeXCfSKP7qD/ACMUAZf9r+MvH/y6FBJ4a0F+upXaZup19Yo/4Qf7x+oPauk8LeAtA8ImSbT7Zpb+XJmv7lvMnlJ5JLHpnuBgGumooAKKKKACiiigAooooAKKKKAOZ+In/JO9f/683/lXjnh4z/D+08M+MrcO2k6lF9m1ONecHc2Gx9ACPdSP4q9j+In/ACTvX/8Arzf+VYvgnRrTxD8GNO0q9TdBc2zofVTvbDD3BwR9K+jy/Exw+Xt1FeEp2kvJx/NbrzRjON56djP+N80Vz8NoJ4ZFkiku4XR1OQylWIIrobyXxe9rp1p4dtdNS3a0jaS+vZGOxscqI15Jxg5PHNeJeIdVvNM8D33gTWWP27Sr+NrZz/y0hIbp7DII9m9q+kdL/wCQRZf9cE/9BFVj6LwOFpRaUlzTtfZpqNn9wRfNJnH/APCu7jVvm8V+JdR1YHrawt9mt/oUTk/XIrp9H8OaNoEXl6VpltaDGC0cYDN9W6n8TWpRXiVcbXqx5ZS93stF9ysjRRSCiiiuUoKKKKACiiigAooooAK5HxL8PNJ1+7Gp2zzaTracx6lYtskz/tgcOPrzjjIrrqKAPN08Y+I/BLrb+ObH7XpwIVNe0+MlAP8AptGOUPuOOwB6132n6jZatYx3un3UN1ayjKSwuGU/iKssqujI6hlYYIIyCK4DUPhxJpl7Jq3gXUf7Dv3O6W0I3WdyfRo/4fqvTsO9AGz4o8B6L4pdLm4jktNTi5g1Gzbyp4iOnzDqPY5/Cuc/4SHxb4C/d+KbZtd0ReBrFjHiaJfWaIf+hD9TV3S/iQLS/j0fxpp7eH9UY4jlkbdaXHvHL0H0PTpnNd4CGUEEEEcEd6AKOj63pniDT0v9Jvoby2fo8TZwfQjqD7Hmr9c5Z+BtC03xQfEGnW72V26Ms0dtIUhmz3eMcEjn8TnrXR0AFFc74o8baH4SiT+0bktdS8QWUC+ZPMewVB/M4HvXMf2Z4z8ffNrE0vhjQG6WNq+bydf+mj/wA+g56gjvQBqa98R7Gy1A6NoNpLr+u9PslmcrEfWWT7qD17jvis+38A6r4ouI7/4gaiLpFO+LRbNilpEe249ZCPf9RXY6D4b0fwxp62OjWENpAPvbB8zn1Zjyx9ya1KAIra2gs7aO3tYY4IIxtSONQqqPQAcCpaKKACiiigAooooAKKKKACiiigAooooAyvE2kvr3hnUdKjlWJ7uBoldhkKT3NQ+EdDk8N+FdP0eWZZpLVCpkQYDZYnp+NbdFbe3n7H2F/dvf52sKyvc89+JHwyj8cS2l5a3MdnfwgxvI6EiSPqAcdwen1Nd5aQm2soICQxijVCR3wMVNRV1cXWq0YUZu8YXt5XBRSdwooormGFFFFABRRRQAUUUUAFFFFABRRRQAUUUUAU9U0nT9asJLHU7OG7tZPvRTIGH19j79a4I+FfFHgUmbwZeHVNJXltD1CXlB6QSn7v0PH1Nek0UAcJp3xd8KXFtdf2ndvo19Zj/SbLUEKSoR1Cj+P6Lk+wrPXxN4r8fAJ4RtTo2ivwdav48ySDv5MX/sx4+hrtNV8K6Brl5b3eqaPZ3lxb/6uSaEMR7c9R7HitcAAAAYA6CgDl/DHgLRvDEr3kSS3urS8z6leN5k8h7/ADH7o9h+Oa6iiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKAP/2Q==",
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASwAAAEsCAIAAAD2HxkiAAAi4klEQVR4nO3deVxU9foH8IdhBxdQLA3JpVTE5WK4Y3nNel2X6WWloqajaTZR1hj3mlhpk+XCNavRNMVsGXfJROcm6g/NBTUlTTMUlxAEKU0EjU2Wmef3x3caCY0Zzpwz33OG5/3yD4HvnPOwfGbOnPN8v8cDEYEQwo+KdwGENHQUQkI4oxASwhmFkBDOKISEcEYhJIQzCiEhnFEICeGMQkgIZxRCQjijEBLCGYWQEM4ohIRwRiEkhDMKISGcUQgJ4YxCSAhnFEJCOKMQEsIZhZAQziiEhHBGISSEMwohIZxRCAnhjEJICGcUQkI4oxASwhmFkBDOKISEcEYhJIQzCiEhnFEICeGMQkgIZxRCQjijEBLCGYWQEM4ohIRwRiEkhDMKISGcUQgJ4YxCSAhnFEJCOKMQEsIZhZAQziiEhHBGISSEMwohIZxRCAnhjEJICGcUQkI4oxASwhmFkBDOKISEcEYhJIQzCiEhnFEICeGMQkgIZxRCQjijEBLCGYWQEM4ohIRwRiEkhDMKISGcUQgJ4YxCSAhnFELCTWpq6ksvvbRr1y7ehXDmgYi8ayANztWrV0eOHHnkyBH2Ybdu3ZKTkx966CG+VfFCr4TEpSwWyxtvvBEaGsoSGBgYCAA///xzx44d4+LizGYz7wI5oBAS10lPT4+Ojl68eLHFYgkMDPzss89KSkq2bt3arFkzi8ViMBi6deu2e/du3mW6HBIivStXrmg0Gg8PDwAIDQ1dsGCB2WyuOWDZsmXt2rVjf5NqtfrSpUu8SnU9BbwnzMwEk8n6/wEDIDq6rsFJSZCdDQAwdCh07y55bcSu8vLypUuXzps3r6SkxN/fX6fTzZ49u1GjRnePrKysXLFixZw5c4qLi318fGJjY+fNm9e4cWPX1+xqvJ8F7NuwAQGs/1q2xJs36xo8ZIh15OrVrqqP/D2TydS2bVv2l6ZWq7Ozs+0+JD8/X6vVqlQqAHjggQcSExNrvWa6H4WFEAB1uroGUwhl4sSJE48++iiLX48ePQ4ePFivh6enp/fv3589vGfPnocPH5aoTjlQXghVKjx69G8HUwi5Kygo0Ol0np6eANC8eXODwVBdXS1gOxaLJSkp6cEHHwQADw+P0aNHX758WfRq5UBJIWzUyPqfnj3x736tFEKOKisrDQZD06ZNAcDb21un092s+82DA0pKSvR6vZ+fHwAEBgbq9fry8nJRqpUPJYXwqacwKsr6/08+ufdgCiEvqampERER7ADyiSeeOHPmjIgbz83N1Wg0bONhYWFGo1HEjXOnpBCOGIG7d1v/36QJXrlyj8EUQtc7d+7csGHDWEI6der07bffSrSj7777rvufp7wHDRr0008/SbQjF1NYCBFx2DDrhzEx9xhMIXSlwsLC+Ph4Hx8fAAgODk5ISKioqJB0j2az2Wg03nfffQCgUqk0Gs21a9ck3aMLKC+EFy+ir6/1M3c/51IIXYNvGGqGPygoyAXhl5TyQoiIM2ZYP9OmDZaU/GUwhdAFxDosrK6utlgsgsuoeRjcsWNH6Q6DpabIEN66ha1aWT/51lt/GUwhlJSIJ0j2798fGRnp/CkWSU8IuYYiQ4iI69ZZP+njgzV/7BRCiYh4qSAnJ2f06NEsNv369XO+NikujbiSUkOIiIMGWT8/cCDaDmoohKJjF83DwsKcv2heWlqq1+v9/f0BICAgID4+vri4WKw6xWoScD0Fh/D8+TtnaNavt36SQiiumu1jvXr1Etw+xpLcpk0bW5JzcnLELZX58ccfa7bLHThwQIq9iEvBIUTE+Hjrl0JDrWdoKIRiEbGR+vjx49F/Tn6JiopKS0sTt9S7CWgc50jZISwtxbZtrV/V6xHvFcIxY/CFFzAlBZV8EtulysrKEhIS2BwiHx8fnU73xx9/CNvUr7/+aktyq1atEhMTXXaIyL4LNmfK398/Pj5e8HchNWWHEBGTk61fDQjAvLzaISwsRC8v62eCgnDCBExOxrIy134DimIymUSZXMtOljRp0sR2suTWrVviluqIWpOJjUajMxdFJKL4ECKiWm0dMGnSPV4Js7LQYMDoaPTwsH7J3x/VajQakcdfhXydPHnyscceY/Hr3Lnzrl27BG/KZDLZVm1Sq9W//PKLiHUKcOzYsX79+tne2R45coRvPbW4QwizstDf3zrLqU2bv31PmJNTO41+fqhWY2Ii/v67xN+DvN24ccN2XrFZs2bOnFfMzMwcOnQo+3MPDw9PSUkRt1TBLBaL0Whs2bIlOzOk0Wh+++033kVZuUMIEfHdd/8y57DuEzO5uZiYiGr1nSNVT0+MjkaDAWXze3ERdtAYFBRkO2gsKioStimWZC8vL9ZHajAYqqqqRC1WBOxqp6+vr6wmRrlJCG/fxk6dHA2hzfXraDSiWo3e3rXTmJ8v3jcgV6mpqV26dLH1mmRkZAjbTlVVVWJiYosWLQDAy8tLq9X+Lu9Di4sXL9q6BR5++OGkpCS+9bhJCBHx//6v3iG0uXHDmkbbhUeVCqOiUK/Hixedq16Wzp8/r1arbV2X//vf/wRvas+ePd26dWObevzxx0+fPi1inZLau3evTCp3nxAi4qhRAkNoU1qKJhNqNHdm8QNgRATq9ZiZKWSDclNUVBQfH8+Ox9j8g9u3bwvblNxeTwSQyWu4AkK4ZQsGB2NwMD73nJ2ReXnYqpV18Nq1Tu20rMyaxiZNaqfx+HGntszL3ZOPrl69KmxTd7+zEpxkOSgsLOT7blYBIeSrvBxNJtRqsUWLO2ls3x51OkxLQ/ldc7q3ffv2/eMf/2CvWgMHDjx58qSw7cj5HKOTap7X7dSpkyvP61IIHVVRgTt34tSpGBLylzS+8QYeO2aW4SVgJi8vz3a1unXr1s5crT527Fjfvn3Zn2nv3r2///57cUuVAy5XOOUbwosXUa3GvDzeddyluhrT0lCnw9BQaxT79s0PCwvTarUmk0k+5+XZlAU2+SggIECv15cJ7RVSRN+JWFzf6yPTEFZUYM+eCIBaLe9S/p7ZjGlpGBeHw4evhD+1bNkyNjY2NTWVYxrvXrFT8JSFuzswRZx8JGes65U1MISEhEg6MUqmIWQLWLRvb2fRe/nIyMjQ6/Xh4eG2NAYHB2s0GpPJ5OKTFrWmLBw6dEjwppQ1F0EKJ06cGDBgAPsJPPLIIxLN/5BjCHfvRpUKvbxQZi1+DmFpjIqKsqUxICBArVYbjUapX0PunrIgePJRzUXsH3nkkfouYu9mTCYTmwnJnoxEnwkpuxBeu2ZdP2bhQt6lOCcrK8tgMERHR7O3UuxwjqVR9PcYFRUVtrcxbPKR4F0od366pEpLS22H5aKvCSCvEFos1ikRAwf+7UL3ipOTk1MrjX5+fmq1OjExUZRLwyaTqX379rbn6aysLGHbUfpKLS4g4qnmmuQVwg8+QABs0cI9Wzfz8vISExPVajW7LgwAnp6e0dHRBoNB2NW2zMzMIUOG2KYs7Ny5U3BtbrBmmcscPXq0T58+7GfVp08f5y/VyCiEx4+jjw96eOD27fYHl5bikiVKfbUsKCgwGo1qtZotX8taWFgar9xzcf+71JyywCYfCT4ZW2sR+x07dgjbToPC2o/uv/9+59uPUD4hLCmxToN4/XWHxr/0EgLglCkSlyWxwsJClkbWAsZ+o1FRUXq9/sKFC/d8CGt3DAkJsbU7Xr9+XfDeXbyIvZspLi62te81atRo9uzZwlbQkEsIJ05EAOzWDR2Z3vXNNwiAvr546pT0lblEaWmpyWTSaDQ17yMdERGh1+sza3SO79mzp2vXruyrgwcP/vnnn4Xtjj2Rs8Zlt7mjAy8XLlywNbJ7e3vHxcXVdwuyCOHmzQiAgYEOzVTIy8PmzREAly2TvjKXKykp+frrr8eOHVvzXu3du3ePi4t78skn2YcdOnQwmUyCd+Gu9zbiKyUlhTUnsYv7u3fvdvyx/EOYlWWdqfD55/YHm83WNX+HDVNM87Qw5eXlJpNJq9WyeQ8eHh7e3t5OTllw77v8cVdWVhYTE8POnXp6ejr+QM4hrKrCvn0RAEeNcmi8Xm9dZVTo+yDlqays3LVrF7twl5ubK2wjDeF+tzJx/vx59jTn+EM4h3DmTATAsDAsLLQ/+OBB9PRElQr37pW+Mplhp08EvAbevYi94CQTBykphPv2WdvTHFlbvbDQupLanDnSVyY/wkKYnp5ec6k/wYvYk3pRTAh//93anjZ/vkPjY2IQAHv3xspKiSuTpfqGMD8/39bb4eQi9qS+lBFCiwWfeqoe7WnLl1uX0G54ffxWjodQxEXsiTDKCOGHHyIABgejI/fYysiwru27caP0lcmVgyEUaxF74gwFhPD0afTzQw8P3LbN/uDycuzeXe6ze13AkRCazeZevXoBQGRkpCJuCaY4xcXF2dnZBQUFdQ+TewhLSjA8HAFQp3NofGysdZmz0lKJK5M3B18JDx8+7Mo7HzU0n3/+OQBMnjy57mH1DaEXuNarr8K5c9C1KyQk2B+cnAwrV4KvL6xfDwEB0henfP3797fd05MohcqVO/v6a/jqKwgIgKQk8Pe3M/jKFXjxRQCAxYshMlL64gjhxHUhvHTJGqolS6BzZzuDLRaYOBFu3IBhw2DaNBdURwg3LgphdTVMmAC3bsHIkTB1qv3x770H+/ZBaCgYjfDnfHRC3JOLQjhnDnz/PYSFwapV9gcfOgTz5oFKBUYjhIRIXxwhXLkihAcOwAcfgJcXbNwIzZrZGXzzJkyYAGYzvPUWDB7sguoI4UzyEF6/fn3x4mQvL9Dr4c/lMOsSGwuXL0Pv3vDOO1KXRogsSBtCRJw8efK33z779NNz33zT/viVK2HzZmjaFDZtAm9vSUsjRC6kDeGSJUt27NgRHBy8aNFkT087g8+ehf/8BwBgxQr4s/WKEPcnYQh//vnnN998EwC+/PJLdl+EOpSXV40dC2Vl8MILMG6cdEURIjtShbC0tDQmJub27duvvvrqiBEj7I6fMeP1gICEqKjqjz+WqCJCZEqqtjWdTnfu3LkuXbosWrTI7uBt27Z9+umnvr6+x44Nbdz4HxKVRIg8SfJKuGXLli+++MLPz2/Dhg3+9vrT8vPzp06dCgCLFi2y3U2WkIZD/BDm5eVptVoAWLJkiW1pvb9jsVgmTpx448aNoUOHvvbaa6IXQ4j8iRzC6urqsWPHFhUVPfvssyyKdZs3b9533313//33f/nllx7Un0YaJJFDqNfrjxw50rp1688++8zu4PT09Hnz5qlUqnXr1rFl/QlpgMQM4YEDB/773/+qVKq1a9c2s9efdvPmzTFjxlRVVc2aNeuJJ54QsQxClEW0EBYVFU2cONFsNuv1+n/+8592x7/88ss5OTm9evV69913xaqBECUSJ4SsPS03N/fRRx99++237Y5ftWrVpk2bmjZtunnzZm/qTyMNmzgh/OSTT7Zv3x4cHLx27VpPe/1pZ8+ejYuLA4BPP/20HfWnkQZPhBBmZGTMmjULAFauXNmmTZu6B1dUVIwfP76srGzKlCnPPfec83snROmcDSG7E015efkrr7wSExNjd/yMGTNOnTr18MMPGwwGJ3dNiHtwNoTTp0/PzMzs0qXL4sWL7Q5OSUlZvny5r69vUlJSzfvvEdKQORXCb775ZvXq1Y63p02aNAkRExISevTo4cx+CXEnwkNoa0/7+OOPHWlPmzRpUkFBwZAhQ6ZPny54p4S4H4EhrK6uHjduXGFh4TPPPBMbG2t3/IIFC/bu3UvtaYTcTWAI586de/jwYcfb09577z3WSdOyZUtheyTEXQkJ4cGDBxcuXKhSqdasWdO8efO6B9+6dWvs2LFVVVUzZ8588sknBRVJiDurdwiLioo0Go3ZbJ4zZ86gQYPsjn/55Zezs7N79uw5d+5cQRUS4ubqHULWnjZgwIA5c+bYHbx69eqNGzc2atRo/fr17L5ChJBa6hfC5cuXb9++PSgoaN26dXbb0y5evPjvf/8bAFauXNmxY0fhNRLi1uoRwm3btrGrCw62p40ZM6a4uPj5558fP368UzUS4tYcDWFZWZlWqzWbzZGRkWPGjLE7Pj09/ezZsw8++OAnn3ziXIWEuDmX3p+QEHI3R0MYEBCwatUqT0/PU6dObd682e743r17R0RE5Obm0vJNhNStHq+ETz/99JIlSwAgNjb28uXLdQ/29fXdvHlz48aNv/rqq/Xr1ztVIyFurX6Ho9OmTRsxYsTNmzcnTJhgNpvrHtyhQ4ePPvoIAGJjYy9cuCC8RkLcWr3fE7IbSxw6dOj999+3O3jq1Knjxo0rKSkZP358ZWWloAoJcXP1DqFtDYv3339/3759dsevWLGiXbt2x48f1+v1giokxM0JOTv62GOPvfnmm7bFs+se3LRp002bNnl7ey9atCg1NVVQkYS4M4GXKPR6fXR09JUrV1588UW7g3v37v3OO+9YLBaNRnP16lVheyTEXQkMoZeX18aNG5s1a5acnLxy5Uq74996663Bgwdfu3Zt8uTJiChsp4S4JeEX68PCwlatWgUAcXFxp0+ftrMblcpoNIaEhOzatYtd5yCEME51zIwcOXLq1Km3b99+7rnnysvL6x4cGhpqNBo9PDxmzZp18uRJZ/ZLiDtxtm1tyZIlnTt3PnPmzIwZM+wOHjZs2LRp0yoqKmJiYoqLi53cNSHuwdkQBgQEJCUl+fv7f/rpp0lJSXbHL168ODIy8pdffnn99ded3DUh7kGEBu6uXbsmJCSAw+1s69evDwgI+OKLLzZs2OD83glROnFmUbz22msjRoywrXxR9+CIiIiPP/4YAF555ZXs7GxRCiBEucQJoYeHB2tnS0tLmz9/vt3xWq127Nixt27dYrcoFKUGQhRKtPmEwcHBa9as8fT0nDt37v79++2OX7FiRdu2bX/44Qe6PyFp4MSc1Dtw4MD4+HjWGVNYWFj34KCgIHZzwoSEhD179ohYBiHKIvLM+rlz5/bv39/xdrbZs2dbLJYJEyZcu3ZN3EoIUQqRQ+jl5bVp06bg4OCtW7eyfpq6zZ49+/HHH6d2NtKQib/GjK2dbfr06Y60s7FlvHfu3ElLQpGGSZKFnkaNGjVlyhTH29lWr14NADNnzvzpp5+kqIcQOZNqtbWlS5eGh4efOXNm5syZdgc//fTTr7zySmTkuy+80IW62UhDI1UIAwMDk5KS/Pz8li1btn37drvjFy82lJXNOnHCKy5OoooIkSkJ1x3t1q3bwoUL4c/bV9Q92N/fe9MmCAiAzz+HjRulK4oQ2ZF28d/p06cPHz68qKho5swv7XWzQUQEfPghAMDLLwN1s5GGQ9oQsnY2tXrrtm36hQvtj4+NhTFj4NYtGDsWqJuNNBCSL4PfokWLGTOeqa6GuXPh8GH741euhDZtID0d3ntP6tIIkQVX3Iti4EB44w2oroZx48BeNxsEBcG6deDpCQsWwN69LqiOEM5cdEOY99+Hfv0gLw+0WvuDBwyA2bPBYoFJk6CgQPriCOHKRSH08oJ166BpU/jmG1i92v74d96BQYMgPx8mTQLqZiPuzXW3RmvfHj77DABg+nTIzLQzWKWCNWugeXNISYHly11QHSHcuPT+hKNHw/PPQ1kZxMSAvW42aN3aGtoZM+DUKemLI4QTV98kdNkyCA+HjAyYNcv+4GeegdhYqKiA8eOhrEz64pTvyJEjq1atsrvCCJEXdLnTp9HPDz08cNs2+4PLy7F7dwRArVb6ymTMx8cHAG7fvl3HGLPZ3LNnTwDo3Lnzrl27XFZbw1FcXJydnV1QUFD3sPomi0MIEfHDDxEAg4Px8mX7gzMy0N8fAXDjRukrkytHQoiIJpOpXbt27I9ArVZfunTJNeWRmpQRQosFn3oKAXDgQKyutj9++XIEwKAgzM6WvDZ5cjCEiFhRUWEwGBo3bgwAPj4+Op3ujz/+cEGFxEYZIUTE33/HVq0QAOfPd2h8TAwCYO/eWFkpcWWy5HgImfz8fK1Wq1KpAOCBBx5ITEw0m82SVkhsFBNCRNy3D1Uq9PLCw4ftDy4sxDZtEADnzJG+MvmpbwiZ9PT0/v37s7+Jnj17HnbkB02cpqQQIuLMmQiAYWFYWGh/8MGD6OmJKhXu3St9ZTIjLISIaLFYkpKSwsLCAMDDw2P06NGXHXkjTpygsBBWVWHfvgiAo0Y5NF6vRwAMDcXr1yWuTDYqKyt37drl6ekJALm5ucI2UlJSotfr/fz8ACAwMFCv15eXl4tbJ2HOnz+vsBAiYlYWNmmCAPj55/YHm804aBAC4LBhaLFIXxw/5eXlJpNJq9Xed9997EXM29ub5UfA6yGTm5ur0WjYn0hYWJjRaBS35gautLQ0JibGw8MDADw9PR1/IP8QIuLmzQiAgYGYmWl/cF4eNm+OALhsmfSVuVxJScnXX389duxYdnqT6d69e1xc3JNPPsk+7NChg8lkEryL7777rnv37mxTgwYN+umnn0Ssv8FKSUlhBxoAEBISsnv3bscfK4sQIuLEiQiA3bqhI0dJ33yDAOjri6dOSV+ZS5SWlppMJo1G06hRI1v2IiIi9Hp9Zo1npj179nTt2pV9dfDgwadPnxa2O7PZbDQa2WusSqXSaDTXrl0T6VtpcC5cuDB69Gj2S/H29o6Li6vvFuQSwpIS7NQJAfD11x0a/9JLCIBTpkhclsQKCwuNRqNarfb19WW/RZVKFRUVpdfrL1y4cM+HVFVVJSYmhoSEAICXl5dWq/39998F7z0+Pp6d8gkODk5ISKioqHDiu2lwiouL9Xo9+901atRo9uzZwi7JyiWEiHj8OPr4oIcHbt9uf3BpKS5Z4tCFfhkqKChg2WMBYNmLjo42GAxXrlxxZAs3btzQ6XReXl4A0KxZM4PBUFVVJayYc+fODRs2jJXRqVOnb7/9Vth2GhR2KHH//ffbDiWuXr0qeGsyCiEifvABAmCLFpifz7sUCeTl5SUmJqrVahYe9vadZe+3334TsMHMzMwhQ4awTYWHh6ekpAiuLTU1NSIigm3qiSeeOHPmjOBNub2jR4/26dOH/az69Olz9OhRJzcorxBaLKhW16OdTRFycnIMBkN0dDQ7bwYAfn5+arU6MTFR8JFkTSaTqX379rZ+0V9++UXYdiorKw0GQ9OmTdl7G51Od/PmTefLcyd5eXkajYb9Hlu3bm00Gi1inKOXVwgR8do1azvbwoW8S3FOVlZWrez5+/ur1Wqj0Xjr1i1x98X6RZs0aWLLj+BdFBQU6HQ6dlmyefPmBoOh2m2eDp1QWlqq1+v9/f0BICAgID4+vri4WKyNyy6EiLh7t7Wd7cgR3qXUX0ZGhl6vj4qKsp3kDAgIYNkT8dd2T7/++qutX7RVq1bO9IueOHHi0UcfZfX36NHj4MGD4paqIKzlqE2bNrZjjZycHHF3IccQIuKMGQiA7dujUg6IWPbCw8Nt2QsODtZoNCaTSfC1dWGOHz8eHR3NaoiKikpLSxO8KZPJ1LZtW9sfX3bDm8Ny/PjxAQMGsJ/AI4884swPsw4yDWFFBfbsKfe5vGYzpqVhXBwOH/6qLXstW7aMjY1NTU0VfLrSeTWfvFm/qOAn77KysoSEBHb10t/fX9zDMDljhxXssJwdVkh3WC7TECLixYuoVmNeHu867lJdjWlpqNNhaCgCIAD26bM1LCxMq9WaTCaO2aul1tsYvV5fVlYmbFNXrlyxnZAIDQ0V64SEPLETVKK8wXaQfEMoNxUVuHMnTp2KISHW7LED5jfewGPHzLL9oxTxhN6xY8f69u3LXvB79ep1RIlv2e0xmUwPPfSQ86ea64VCaEd5OZpMqNViixZ/yZ5Oh2lpimki379/f2RkJPvbGjhw4MmTJ4Vtx2KxGI3Gli1bsgNdjUYj7AqnDGVmZg4dOlSUi671pYAQbtmCwcHWf9Om2Rk8apR15Nq1Tu20rAxNJtRorDM82L+ICNTr8fhxp7bMi4hNHmxiFGvXcoOJUYWFhbb2o+DgYGfaj4RRQAg3bLgTA5XKzjT8IUOsI1evFrKv0lJr9ho1qp09R2Z4yF9RUVF8fDzLT1BQUEJCguCTtxcvXrQ1Lj/88MNJSUniluoCrBG3RYsWzjfiOkNhIQTArl3rWmZGWAhv3ECjEdVq9PW9k/aoKNTr8eJF578D2Tl//rxarWb56dChgzP52bt3b7du3dimHn/8ccETO1xPPpUrL4QAuGjR3w6uVwivX7dmz9vb+ihPT4yORoPBPZtXa0lNTe3SpYutXzQjI0PYdmTyeuI4ub2GKymEXbta0xIYiH933cuREObmYmIiqtXo5VU7e+5ylsFR7HR8UFAQOx2v1WqvC104hPs7K0fc/W7Wxa0U96SkEI4Yga++av3/0KH3HlxHCHNy0GDA6Gj08LCO8fNDtRoTE1HeT9ySYxOj2IVpNjFK8IVpjucY6ybn87oKC+G1a3dOVyYn32Pw3SHMyqqdPX9/VKvRaESJr8EqzNmzZ//1r3+x/Di5kD6Xq211qHmFs3fv3nK7wqmwEOKfcw7ZQol3d1DVCmFh4Z1jzqAgnDABk5NRaN9Ig1BrIf2srCxh23F938k9KaLXR3khrKiwLoQBgP/5T+3Bd78SjhmDL7yAKSlISzc4SMSF9F3ZgVmLgrpelRdCRNyxw/oZL6/aaz05eZ2Q2Ii4kL5r5iLUpKz5H4oMIaL1fjIA2KsX1vzboBCKS6yF9F0wK49R4kxIpYbw8mUMDLR+fsWKO5+nEIqO5efBBx8EpxfSLy0ttR0iij4/XblrAig1hIg4f77188HBaFs1k0IoEREX0s/JybFdK+/Xr5/ztSl9dRwFh7CiAsPDrV+yzf2lEEpKxIX02Zpla9ascbIkN1gnTsEhRMTU1Dutnj/8gEghdAmxFtKvrq525oJBrRVTd+zYIXhTfCk7hIg4Zoz1qwMHIlIIXYXvQvputna44kP422/YtKl1wNatFEKXqhkGNjFK6jCw8LNmcbe5i4biQ4iIH35oHdCxIw4eTCF0tXPnzg0fPpwdFnbs2FG6hfTd9X5S7hDCqirs3t06xjYpiULoYpKeIHHvOyu6QwgR8dChO/3ZFEJepLhU0BDuMewmIUTEyZMphLIg1kVz1iQQFhbmfJOAzLlPCAsK/rIYIYWQrx9//LFm+9iBAwfq9fD09PR+/fqxh/fq1Utwu5wiuE8IEXHFCgqhvAhopM7Pz7dNPnKycVwpPBAR5C0zE0wmAIAOHeDZZ+saabHA0qVQUQEAMHQo/HkijfBUXl6+dOnSefPmlZSU+Pv763S6t99+m82TuufI+fPnFxcX+/j4xMbGzps3754j3Q3vZwHSINSaXLtgwYJar2/Lli2rOZn40qVLvEp1PQohcZ1jx47Z3ukFBgauWrUKETdu3MhOqAJAZGTk/v37eZfpago4HCXuxGKxxMfHf/TRRxaLBQACAwNLS0sBQKVS6XS6xYsXs9OqDQqFkHBw9erVkSNHHjlyhH3Yo0eP5ORk25TfhoZCSLhJTU3dsmXLM888M2TIEN618EQhJIQzFe8CCGnoKISEcEYhJIQzCiEhnFEICeGMQkgIZxRCQjijEBLCGYWQEM4ohIRwRiEkhDMKISGcUQgJ4YxCSAhnFEJCOKMQEsIZhZAQziiEhHBGISSEMwohIZxRCAnhjEJICGcUQkI4oxASwhmFkBDOKISEcEYhJIQzCiEhnFEICeGMQkgIZxRCQjijEBLCGYWQEM4ohIRwRiEkhDMKISGcUQgJ4YxCSAhnFEJCOKMQEsIZhZAQziiEhHBGISSEMwohIZxRCAnhjEJICGcUQkI4oxASwhmFkBDOKISEcEYhJIQzCiEhnFEICeGMQkgIZxRCQjijEBLCGYWQEM4ohIRwRiEkhDMKISGcUQgJ4YxCSAhnFEJCOKMQEsIZhZAQzv4f6jkImFkQ6PQAAAAASUVORK5CYII=",
+      "text/plain": [
+       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=300x300>"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Generate Mol Viz\n",
+    "from rdkit import Chem\n",
+    "from rdkit.Chem import Draw\n",
+    "\n",
+    "input_ids = tokenizer(\"<s>\", return_tensors=\"pt\").input_ids.to(device)\n",
+    "gen = model.generate(input_ids, max_length=256, top_k=50, temperature=1, do_sample=True, pad_token_id=tokenizer.pad_token_id)\n",
+    "generatedmol = tokenizer.decode(gen[0], skip_special_tokens=True)\n",
+    "\n",
+    "test = generatedmol.replace(' ', '')\n",
+    "csmi_gen = sf.decoder(test)\n",
+    "print(csmi_gen)\n",
+    "mol = Chem.MolFromSmiles(csmi_gen)\n",
+    "\n",
+    "# Draw the molecule\n",
+    "Draw.MolToImage(mol)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

train-withmtp.py

@@ -0,0 +1,365 @@
+# ========================
+#  Train with NTP + MTP
+#  by gbyuvd
+# ========================
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import json
+import os
+import math
+from typing import List, Union, Optional, Tuple, Dict, Any
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers import Qwen3Config, Qwen3ForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from transformers.models.qwen2.modeling_qwen2 import Qwen2PreTrainedModel
+from datasets import load_dataset, DatasetDict
+import pandas as pd
+from torch.utils.data import Dataset, DataLoader, random_split
+from sklearn.model_selection import train_test_split
+from ranger21 import Ranger21
+from tqdm.notebook import tqdm
+from FastChemTokenizerHF import FastChemTokenizerSelfies
+from ChemQ3MTP import ChemQ3MTP
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+from transformers import TrainerCallback
+import datetime
+
+# ==============================
+# Load external configuration
+# ==============================
+with open("config.json", "r") as f:
+    CONFIG = json.load(f)
+
+TRAINING_CFG = CONFIG["training"]
+MODEL_CFG = CONFIG["model"]
+GENERATION_CFG = CONFIG.get("generation", {})
+
+# Training params
+BATCH_SIZE = TRAINING_CFG["batch_size"]
+NUM_EPOCHS = TRAINING_CFG["num_epochs"]
+LEARNING_RATE = TRAINING_CFG["learning_rate"]
+WEIGHT_DECAY = TRAINING_CFG["weight_decay"]
+GRAD_ACCUM_STEPS = TRAINING_CFG["gradient_accumulation_steps"]
+TOKENIZE_BATCH_SIZE = TRAINING_CFG["tokenize_batch_size"]
+TRAIN_SPLIT_RATIO = TRAINING_CFG["train_split_ratio"]
+VAL_SPLIT_RATIO = TRAINING_CFG["val_split_ratio"]
+TEST_SPLIT_RATIO = TRAINING_CFG["test_split_ratio"]
+INCLUDE_FOR_METRICS = TRAINING_CFG.get("include_for_metrics", ["input_ids", "attention_mask", "labels"])
+# ==============================
+
+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")
+
+
+def main():
+    # --- Load the tokenizer ---
+    tokenizer = FastChemTokenizerSelfies.from_pretrained("./selftok_core")
+
+    out = tokenizer("[C] [=C] [Branch1]", return_tensors="pt")
+    print(out.input_ids)
+    print(out.attention_mask)
+    out = out.to("cuda" if torch.cuda.is_available() else "cpu")
+    print(out.input_ids.device)
+
+    # --- Define config ---
+    config = Qwen3Config(
+        vocab_size=len(tokenizer),
+        bos_token_id=tokenizer.bos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+        tie_word_embeddings=True,
+        use_cache=False,
+        **MODEL_CFG
+    )
+
+    model = ChemQ3MTP(config, num_future_tokens=3)
+
+    def count_parameters(model):
+        return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print(f"Enhanced model has {count_parameters(model):,} trainable parameters.")
+
+    batch_size, seq_len = 2, 32
+    dummy_input = torch.randint(
+        low=0,
+        high=len(tokenizer),
+        size=(batch_size, seq_len),
+        dtype=torch.long,
+    )
+    with torch.no_grad():
+        outputs = model(dummy_input)
+        logits = outputs.logits
+    print(f"Input shape: {dummy_input.shape}")
+    print(f"Logits shape: {logits.shape}")
+
+    print("Loading dataset...")
+    dataset = load_dataset(
+        'csv',
+        data_files='./data/sample_all_14k.csv',
+        split='train',
+        streaming=True
+    )
+
+    print("Shuffling and splitting dataset...")
+    shuffled_dataset = dataset.shuffle(seed=42, buffer_size=10000)
+
+    total_lines = 14000
+    test_size = int(TEST_SPLIT_RATIO * total_lines)
+    val_size = int(VAL_SPLIT_RATIO * total_lines)
+    train_size = total_lines - test_size - val_size
+
+    test_dataset = shuffled_dataset.take(test_size)
+    remaining = shuffled_dataset.skip(test_size)
+    val_dataset = remaining.take(val_size)
+    train_dataset = remaining.skip(val_size)
+
+    print(f"Dataset split: train={train_size}, val={val_size}, test={test_size}")
+
+    def tokenize_function(examples):
+        batch_results = {"input_ids": [], "attention_mask": [], "labels": []}
+        smiles_list = examples['SELFIES'] if isinstance(examples['SELFIES'], list) else [examples['SELFIES']]
+        for smiles in smiles_list:
+            tokenized = tokenizer(
+                smiles,
+                truncation=True,
+                padding=False,
+                max_length=MODEL_CFG["max_position_embeddings"],
+                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
+
+    print("Tokenizing datasets...")
+    train_dataset = train_dataset.map(tokenize_function, batched=True, batch_size=TOKENIZE_BATCH_SIZE, remove_columns=["SELFIES"])
+    val_dataset = val_dataset.map(tokenize_function, batched=True, batch_size=TOKENIZE_BATCH_SIZE, remove_columns=["SELFIES"])
+
+    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 = {key: torch.tensor(values, dtype=torch.long) for key, values in batch.items()}
+            return batch
+
+    data_collator = EnhancedDataCollator(tokenizer, pad_to_multiple_of=8)
+
+    def create_enhanced_optimizer(model_params):
+        num_batches_per_epoch = train_size // BATCH_SIZE
+        optimizer_params = {
+            'lr': LEARNING_RATE,
+            'weight_decay': WEIGHT_DECAY,
+            'use_adabelief': True,
+            'use_cheb': False,
+            'use_warmup': True,
+            'use_madgrad': True,
+            'num_epochs': NUM_EPOCHS,
+            'using_gc': True,
+            'warmdown_active': True,
+            'num_batches_per_epoch': num_batches_per_epoch
+        }
+        return Ranger21(model_params, **optimizer_params)
+
+    from torch.optim.lr_scheduler import LambdaLR
+    class EnhancedCustomTrainer(Trainer):
+        def create_optimizer(self):
+            self.optimizer = create_enhanced_optimizer(self.model.parameters())
+            return self.optimizer
+        def create_scheduler(self, num_training_steps, optimizer=None):
+            if optimizer is None:
+                optimizer = self.optimizer
+            self.lr_scheduler = LambdaLR(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
+
+    steps_per_epoch = train_size // BATCH_SIZE
+    total_steps = steps_per_epoch * NUM_EPOCHS
+
+    training_args = TrainingArguments(
+        output_dir='./chemq3minipret',
+        max_steps=total_steps,
+        per_device_train_batch_size=BATCH_SIZE,
+        per_device_eval_batch_size=BATCH_SIZE,
+        gradient_accumulation_steps=GRAD_ACCUM_STEPS,
+        logging_dir='./gptlo-1',
+        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,
+    )
+
+    print("Initializing enhanced trainer with MTP capabilities...")
+    trainer = EnhancedCustomTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+        data_collator=data_collator,
+        processing_class=tokenizer,
+        callbacks=[LossLoggerCallback("training_losses.txt", with_timestamp=True)]
+    )
+
+    model.set_mtp_training(True)
+    print(" MTP training mode enabled")
+
+    print("Starting enhanced training with MTP and Horizon Loss...")
+    try:
+        print("\n Phase 1: Warmup with standard Causal LM...")
+        model.set_mtp_training(False)
+        warmup_steps = max(1, total_steps // 5)
+        trainer.args.max_steps = warmup_steps
+        trainer.train()
+        print("\n Phase 2: Full MTP + Horizon Loss training...")
+        model.set_mtp_training(True)
+        trainer.args.max_steps = total_steps
+        trainer.train(resume_from_checkpoint=True)
+        print("Enhanced training completed successfully!")
+        trainer.save_model("./enhanced-qwen3-final")
+        tokenizer.save_pretrained("./enhanced-qwen3-final")
+        training_config = {
+            "model_type": "EnhancedQwen3ForCausalLM",
+            "num_future_tokens": 3,
+            "horizon_loss_enabled": True,
+            "mtp_head_enabled": True,
+            "training_phases": ["causal_lm_warmup", "mtp_horizon_training"],
+            "total_parameters": count_parameters(model),
+        }
+        config_path = "./enhanced-qwen3-final/training_config.json"
+        with open(config_path, "w") as f:
+            json.dump(training_config, f, indent=2)
+        print(f" Enhanced model, tokenizer, and config saved!")
+    except Exception as e:
+        print(f"Enhanced training failed with error: {e}")
+        import traceback
+        traceback.print_exc()
+        return
+
+    print("\nmTesting enhanced generation capabilities...")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    model.eval()
+    try:
+        print("\n--- Standard Generation Test ---")
+        input_ids = tokenizer("<s> [C]", return_tensors="pt").input_ids.to(device)
+        with torch.no_grad():
+            model.set_mtp_training(False)
+            gen = model.generate(
+                input_ids,
+                max_length=GENERATION_CFG.get("max_length", 64),
+                top_k=GENERATION_CFG.get("top_k", 50),
+                top_p=GENERATION_CFG.get("top_p", 0.9),
+                temperature=GENERATION_CFG.get("temperature", 0.8),
+                do_sample=GENERATION_CFG.get("do_sample", True),
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                num_return_sequences=GENERATION_CFG.get("num_return_sequences", 3),
+            )
+            for i, sequence in enumerate(gen):
+                result = tokenizer.decode(sequence, skip_special_tokens=True)
+                print(f"Generated SELFIES {i+1}: {result}")
+        print("\n--- MTP Analysis Test ---")
+        model.set_mtp_training(True)
+        test_smiles = "[C]"
+        test_input = tokenizer(test_smiles, return_tensors="pt", add_special_tokens=True).to(device)
+        with torch.no_grad():
+            outputs = model(**test_input)
+            if hasattr(model.mtp_head, 'prediction_heads'):
+                hidden_states = model.model(test_input['input_ids']).last_hidden_state
+                mtp_outputs = model.mtp_head(hidden_states)
+                print(f"Input SELFIES: {test_smiles}")
+                print(f"Tokenized: {tokenizer.convert_ids_to_tokens(test_input['input_ids'][0].tolist())}")
+                for i, (key, logits) in enumerate(mtp_outputs.items()):
+                    top_tokens = torch.topk(logits[0], k=3, dim=-1)
+                    print(f"\n{key} predictions:")
+                    for pos in range(min(5, logits.size(1))):
+                        pos_preds = []
+                        for j in range(3):
+                            token_id = top_tokens.indices[pos, j].item()
+                            prob = torch.softmax(logits[0, pos], dim=-1)[token_id].item()
+                            token = tokenizer.id_to_token.get(token_id, '<UNK>')
+                            pos_preds.append(f"{token}({prob:.3f})")
+                        print(f"  Position {pos}: {', '.join(pos_preds)}")
+        print("\nEnhanced generation tests completed!")
+    except Exception as e:
+        print(f"Enhanced generation test failed: {e}")
+        import traceback
+        traceback.print_exc()
+
+    print("\nEnhanced Model Analysis:")
+    print(f"Total parameters: {count_parameters(model):,}")
+    mtp_params = sum(p.numel() for p in model.mtp_head.parameters() if p.requires_grad)
+    horizon_params = sum(p.numel() for p in model.horizon_loss.parameters() if p.requires_grad)
+    base_params = count_parameters(model) - mtp_params - horizon_params
+    print(f"Base model parameters: {base_params:,}")
+    print(f"MTP head parameters: {mtp_params:,}")
+    print(f"Horizon loss parameters: {horizon_params:,}")
+    print(f"Enhancement overhead: {((mtp_params + horizon_params) / base_params * 100):.2f}%")
+    print(f"\n Enhanced Model Architecture:")
+    print(f"- Base Model: Qwen3 with {config.num_hidden_layers} layers")
+    print(f"- Hidden Size: {config.hidden_size}")
+    print(f"- Attention Heads: {config.num_attention_heads}")
+    print(f"- Vocab Size: {config.vocab_size}")
+    print(f"- MTP Future Tokens: {model.mtp_head.num_future_tokens}")
+    print(f"- Horizon Loss Weights: Learnable")
+    print(f"- Training Mode: {'MTP + Horizon Loss' if model.use_mtp_training else 'Standard Causal LM'}")
+    print("\n Enhanced training pipeline completed successfully!")
+
+if __name__ == "__main__":
+    main()

train_ppokl_withsa.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+# Refactored PPO-KL training script using ChemQ3MTP module
+
+import os
+import torch
+from tqdm import tqdm
+from FastChemTokenizerHF import FastChemTokenizerSelfies
+from ChemQ3MTP import ChemQ3MTP, CurriculumManager
+
+
+def main():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"🚀 Using device: {device}")
+
+    # --- Load tokenizer ---
+    tokenizer = FastChemTokenizerSelfies.from_pretrained("../selftok_core")
+
+    # --- Load model ---
+    model = ChemQ3MTP.from_pretrained("../pretrained/sample-e1-mtp")
+    model.tokenizer = tokenizer
+    model.to(device)
+
+    # --- RL fine-tuning setup ---
+    print("\n🎯 Phase 2: RL Fine-tuning with PPO + Curriculum Learning")
+    model.set_mtp_training(False)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=5e-6)
+    curriculum = CurriculumManager(start_len=10, max_len=35, step_increase=5, steps_per_level=70)
+    baseline = None
+    gamma = 0.95
+
+    # Dummy input (BOS-only batch)
+    batch_size = 4
+    dummy_input = tokenizer([tokenizer.bos_token] * batch_size, return_tensors="pt", padding=True)
+    input_ids = dummy_input.input_ids.to(device)
+
+    # Training config
+    total_steps = 14000
+    checkpoint_steps = {total_steps // 4, total_steps // 2, 3 * total_steps // 4, total_steps}
+    checkpoint_dir = "./ppo_checkpoints"
+    os.makedirs(checkpoint_dir, exist_ok=True)
+
+    # --- RL Training Loop with tqdm ---
+    for step in tqdm(range(total_steps), desc="RL Training"):
+        max_new_tokens = curriculum.get_max_new_tokens()
+
+        # === PPO Rollout ===
+        with torch.no_grad():
+            selfies_list, old_log_probs, _, old_action_probs = model.generate_with_logprobs(
+                input_ids=input_ids,
+                max_new_tokens=max_new_tokens,
+                temperature=1.0,
+                top_k=50,
+                top_p=0.95,
+                do_sample=True,
+                return_probs=True
+            )
+            old_log_probs = old_log_probs.detach()
+            old_action_probs = old_action_probs.detach()
+
+        # === PPO Update ===
+        ppo_result = model.ppo_step(
+            input_ids=input_ids,
+            old_log_probs=old_log_probs,
+            old_action_probs=old_action_probs,
+            tokenizer=tokenizer,
+            max_new_tokens=max_new_tokens,
+            # validity_weight=1.0,   # only used in ChemQ3 mode
+            # lipinski_weight=1.0,   # only used in ChemQ3 mode
+            entropy_weight=0.01,
+            clip_epsilon=0.2,
+            baseline=baseline,
+            reward_mode="sa",        # 🔑 SA-only mode
+        )
+
+
+
+        loss = ppo_result['loss']
+        optimizer.zero_grad(set_to_none=True)  # slightly more efficient than zeroing
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+
+        # === Update baseline ===
+        reward_tensor = torch.tensor(ppo_result['avg_reward'], device=device)
+        baseline = reward_tensor if baseline is None else gamma * baseline + (1 - gamma) * reward_tensor
+
+        # Curriculum update
+        curriculum.step()
+
+        # Checkpointing
+        if (step + 1) in checkpoint_steps:
+            checkpoint_path = os.path.join(checkpoint_dir, f"model_step_{step+1}")
+            model.save_pretrained(checkpoint_path)
+            tokenizer.save_pretrained(checkpoint_path)
+            torch.save({
+                'step': step + 1,
+                'optimizer_state_dict': optimizer.state_dict(),
+                'baseline': baseline.item(),
+                'curriculum_state': {
+                    'current_max_len': curriculum.current_max_len,
+                    'step_counter': curriculum.step_counter
+                }
+            }, os.path.join(checkpoint_path, 'training_state.pt'))
+            print(f"\n💾 Checkpoint saved at step {step+1} -> {checkpoint_path}")
+
+         # Logging every 50 steps
+        if step % 50 == 0:
+            print(f"\n[RL Step {step}] "
+                  f"Loss={loss.item():.4f} | "
+                  f"Valid={ppo_result['validity_rate']:.3f} | "
+                  f"Lipinski={ppo_result['lipinski_score']:.3f} | "
+                  f"Reward={ppo_result['avg_reward']:.3f} | "
+                  f"Entropy={ppo_result['entropy']:.3f} | "
+                  f"EntropyW={ppo_result['entropy_weight']:.4f}")
+
+            sample_selfies = ppo_result['generated_selfies'][0][:100]
+            sample_smiles = ppo_result['generated_smiles'][0] or "Invalid"
+            print(f"  Sample SELFIES: {sample_selfies}")
+            print(f"  Sample SMILES: {sample_smiles}")
+
+
+            sample_selfies = ppo_result['generated_selfies'][0][:100]
+            sample_smiles = ppo_result['generated_smiles'][0] or "Invalid"
+            print(f"  Sample SELFIES: {sample_selfies}")
+            print(f"  Sample SMILES: {sample_smiles}")
+
+
+    print("🎉 Training complete!")
+
+if __name__ == "__main__":
+    main()