Upload x_transformer_2_3_1.py

x_transformer_2_3_1.py

@@ -4,7 +4,7 @@
 #
 # Partial x-transformers code With useful modifications as a stand-alone Python module
 #
-# Version 3.0
+# Version 6.0
 #
 # Original source code courtesy of lucidrains
 # https://github.com/lucidrains/x-transformers
@@ -45,6 +45,7 @@ import torch
 from torch.nn import Module
 from torch import nn, einsum, Tensor
 import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
 
 from collections import namedtuple
 from functools import wraps
@@ -3982,7 +3983,7 @@ class AutoregressiveWrapper(Module):
         # whether to add router z-loss
         self.add_attn_z_loss = add_attn_z_loss
 
-    @torch.no_grad()
+    @torch.inference_mode()
     @eval_decorator
     def generate(
         self,
@@ -4147,7 +4148,7 @@ class AutoregressiveWrapper(Module):
 
         return out
 
-    @torch.no_grad()
+    @torch.inference_mode()
     @eval_decorator
     def generate_masked(
         self,
@@ -4328,7 +4329,7 @@ class AutoregressiveWrapper(Module):
 
         return out
 
-    @torch.no_grad()
+    @torch.inference_mode()
     @eval_decorator
     def generate_biased(
         self,
@@ -4556,6 +4557,226 @@ class AutoregressiveWrapper(Module):
         out, = unpack(out, ps, '* n')
     
         return out
+
+    @torch.inference_mode()
+    @eval_decorator
+    def generate_advanced(
+        self,
+        prompts,
+        seq_len,
+        eos_token = None,
+        temperature = 1.,
+        prompt_lens: Tensor | None = None,
+        filter_logits_fn: str | Callable = top_k,
+        restrict_to_max_seq_len = True,
+        amateur_model: Module | Tuple[Module] | None = None,
+        filter_kwargs: dict = dict(),
+        contrastive_decode_kwargs: dict | Tuple[dict] = dict(
+            beta = 0.5,
+            alpha = 0.1
+        ),
+        cache_kv = True,
+        return_prime=False,
+        verbose=True,
+        # --- new generation options ---
+        logits_bias: dict | None = None,      # {token_id: bias_value} where bias_value is float or Tensor(batch,)
+        masked_tokens: list | Tensor | None = None,  # list of token ids to forbid
+        # --- binary classifier mode ---
+        binary_classifier: bool = False,     # if True, run classifier snippet and return preds, probs
+        classifier_model: Module | None = None,  # model to use for binary classification
+        batches: list | None = None,         # iterable of input batches for classifier_model
+        threshold: float = 0.5,              # threshold for converting probs to preds
+        classifier_device: torch.device | None = None,
+        # -----------------
+        **kwargs
+    ):
+        # If binary classifier mode requested, run the provided snippet and return early.
+        if binary_classifier:
+            assert classifier_model is not None, "classifier_model must be provided when binary_classifier=True"
+            assert batches is not None, "batches (iterable of input tensors) must be provided when binary_classifier=True"
+
+            device = classifier_device if classifier_device is not None else (prompts.device if exists(prompts) else torch.device('cpu'))
+
+            all_probs = []
+            all_preds = []
+
+            classifier_model.eval()
+            with torch.no_grad():
+                for x in batches:
+                    x = x.to(device)
+                    logits = classifier_model(x).squeeze()    # [B]
+                    probs = torch.sigmoid(logits)    # [B]
+                    preds = (probs >= threshold).long()
+
+                    all_probs.extend(probs.cpu().tolist())
+                    all_preds.extend(preds.cpu().tolist())
+
+            return all_preds, all_probs
+
+        # --- normal generation path below ---
+        max_seq_len, greedy, device = self.max_seq_len, temperature == 0., prompts.device
+
+        prompts, ps = pack([prompts], '* n')
+
+        b, t = prompts.shape
+
+        # handle filter logits fn given as string
+        if isinstance(filter_logits_fn, str):
+            assert filter_logits_fn in FILTER_LOGITS_FN, f"only {join(FILTER_LOGITS_FN.keys())} are available"
+            filter_logits_fn = FILTER_LOGITS_FN[filter_logits_fn]
+
+        # handle variable lengthed prompts (prefixes)
+        seq_start_pos = None
+        if exists(prompt_lens):
+            prompts = align_right(prompts, prompt_lens, pad_id = self.pad_value)
+            seq_start_pos = t - prompt_lens
+
+        # output from which sampled tokens appended to
+        out = prompts
+
+        if verbose:
+            print("Generating sequence of max length:", seq_len)
+
+        # kv caches
+        cache = None
+
+        # if doing contrastive decoding, turn off filter automatically
+        if exists(amateur_model):
+            amateur_model = cast_tuple(amateur_model)
+            contrastive_decode_kwargs = cast_tuple(contrastive_decode_kwargs)
+
+            assert len(amateur_model) == len(contrastive_decode_kwargs)
+
+            amateur_caches = [None] * len(amateur_model)
+            filter_logits_fn = identity
+
+            for i, module in enumerate(amateur_model):
+                if isinstance(module, AutoregressiveWrapper):
+                    amateur_model[i] = module.net
+
+                module.eval()
+
+        # normalize inputs for new args
+        if exists(logits_bias):
+            assert isinstance(logits_bias, dict), "logits_bias must be a dict {token_id: bias_value}"
+        if exists(masked_tokens):
+            if isinstance(masked_tokens, torch.Tensor):
+                masked_tokens = masked_tokens.tolist()
+            else:
+                masked_tokens = list(masked_tokens)
+
+        # sampling up to seq_len
+        for sl in range(seq_len):
+
+            if restrict_to_max_seq_len:
+                max_len_exceeded = out.shape[-1] > max_seq_len
+
+                assert not (cache_kv and max_len_exceeded and not self.net.can_cache_kv_outside_max_seq_len), 'the network cannot use cached key values when decoding outside the max sequence length. most likely because you are using absolute positional embedding. you can switch to rotary embeddings to resolve this issue'
+
+                x = out[:, -max_seq_len:]
+
+                if exists(cache):
+                    for inter in cache.attn_intermediates:
+                        if inter.layer_type == 'a':
+                            inter.cached_kv = [t[..., -(max_seq_len - 1):, :] for t in inter.cached_kv]
+
+            logits, new_cache = self.net(
+                x,
+                return_intermediates = True,
+                cache = cache,
+                seq_start_pos = seq_start_pos,
+                **kwargs
+            )
+
+            if cache_kv and self.net.can_cache_kv:
+                cache = new_cache
+
+            logits = logits[:, -1]   # shape: (batch, vocab)
+
+            # handle contrastive decoding, Li et al.
+            if exists(amateur_model):
+                for i, (amateur, amateur_cache, amateur_contrastive_decode_kwargs) in enumerate(zip(amateur_model, amateur_caches, contrastive_decode_kwargs)):
+                    amateur_logits, next_amateur_cache = amateur(
+                        x,
+                        return_intermediates = True,
+                        cache = amateur_cache,
+                        seq_start_pos = seq_start_pos,
+                        **kwargs
+                    )
+
+                    amateur_logits = amateur_logits[:, -1]
+
+                    assert amateur_logits.shape == logits.shape, 'logits dimension are not the same between amateur and expert model'
+                    logits = contrastive_decode_fn(logits, amateur_logits, **amateur_contrastive_decode_kwargs)
+
+                    if cache_kv and amateur.can_cache_kv:
+                        amateur_caches[i] = next_amateur_cache
+
+            # --- APPLY LOGITS BIAS AND MASKING HERE (before filtering / softmax) ---
+            # logits_bias: dict {token_id: bias_value} where bias_value is float or Tensor(batch,)
+            if exists(logits_bias):
+                # apply per-token bias updates directly to logits to avoid allocating full vocab bias tensor
+                for tok_id, bias_val in logits_bias.items():
+                    # support scalar or per-batch tensor
+                    if isinstance(bias_val, torch.Tensor):
+                        if bias_val.dim() == 1 and bias_val.shape[0] == b:
+                            bias_to_add = bias_val.to(device)
+                        else:
+                            bias_to_add = bias_val.to(device).view(1).expand(b)
+                    else:
+                        bias_to_add = torch.tensor(float(bias_val), device=device).view(1).expand(b)
+
+                    logits[:, int(tok_id)] = logits[:, int(tok_id)] + bias_to_add
+
+            # masked_tokens: list of token ids to forbid
+            if exists(masked_tokens) and len(masked_tokens) > 0:
+                NEG_INF = -1e9
+                idx = torch.tensor(masked_tokens, device=device, dtype=torch.long)
+                idx = idx[(idx >= 0) & (idx < logits.shape[-1])]
+                if idx.numel() > 0:
+                    logits.index_fill_(dim=-1, index=idx, value=NEG_INF)
+            # -------------------------------------------------------------------
+
+            # filter by top_k, top_p (nucleus), top_a, or custom
+            if greedy:
+                sample = logits.argmax(dim = -1, keepdim = True)
+            else:
+                filtered_logits = filter_logits_fn(logits, **filter_kwargs)
+                probs = F.softmax(filtered_logits / temperature, dim=-1)
+                sample = torch.multinomial(probs, 1)
+
+            # concat sample
+            out = torch.cat((out, sample), dim=-1)
+
+            if verbose:
+                if sl % 32 == 0:
+                    print(sl, '/', seq_len)
+
+            if not exists(eos_token):
+                continue
+
+            is_eos_tokens = (out == eos_token)
+
+            if is_eos_tokens.any(dim = -1).all():
+                if verbose:
+                    print('Model called the end of sequence at:', sl, '/', seq_len)
+                break
+
+        if exists(eos_token):
+            # mask out everything after the eos tokens
+            shifted_is_eos_tokens = F.pad(is_eos_tokens, (1, -1))
+            mask = shifted_is_eos_tokens.float().cumsum(dim = -1) >= 1
+            out = out.masked_fill(mask, self.pad_value)
+
+        if return_prime:
+            out = out[:, :]
+
+        else:
+            out = out[:, t:]
+
+        out, = unpack(out, ps, '* n')
+
+        return out
     
     def compute_accuracy(self, logits, labels):
         
@@ -4613,6 +4834,438 @@ class AutoregressiveWrapper(Module):
 
         return loss, acc, logits, cache
 
+    @torch.inference_mode()
+    @eval_decorator
+    def generate_expert(
+        self,
+        prompts,
+        seq_len,
+        eos_token = None,
+        temperature = 1.,
+        prompt_lens: Tensor | None = None,
+        filter_logits_fn: str | Callable = top_k,
+        restrict_to_max_seq_len = True,
+        amateur_model: Module | Tuple[Module] | None = None,
+        filter_kwargs: dict = dict(),
+        contrastive_decode_kwargs: dict | Tuple[dict] = dict(
+            beta = 0.5,
+            alpha = 0.1
+        ),
+        cache_kv = True,
+        return_prime=False,
+        verbose=True,
+        # --- new controls ---
+        token_type_ids: torch.LongTensor | None = None,   # [vocab]
+        type_temperatures: dict | None = None,            # {type_id: temp}
+        type_biases: dict | None = None,                  # {type_id: bias}
+        repetition_window: int = 64,
+        repetition_penalty_per_type: dict | None = None,  # {type_id: penalty_scale}
+        rare_types: set | None = None,                    # e.g. {4, 5}
+        rare_type_boost: float = 0.0,                     # small, e.g. 0.5
+        entropy_threshold: float = 2.0,                   # when below, boost rare types
+        # --- masked tokens option ---
+        forbidden_token_ids: torch.LongTensor | torch.BoolTensor | None = None,
+        forbidden_value: float = -1e9,
+        **kwargs
+    ):
+        max_seq_len, greedy, device = self.max_seq_len, temperature == 0., prompts.device
+    
+        prompts, ps = pack([prompts], '* n')
+    
+        b, t = prompts.shape
+    
+        # handle filter logits fn given as string
+    
+        if isinstance(filter_logits_fn, str):
+            assert filter_logits_fn in FILTER_LOGITS_FN, f"only {join(FILTER_LOGITS_FN.keys())} are available"
+            filter_logits_fn = FILTER_LOGITS_FN[filter_logits_fn]
+    
+        # handle variable lengthed prompts (prefixes)
+    
+        seq_start_pos = None
+        if exists(prompt_lens):
+            prompts = align_right(prompts, prompt_lens, pad_id = self.pad_value)
+            seq_start_pos = t - prompt_lens
+    
+        # output from which sampled tokens appended to
+    
+        out = prompts
+    
+        if verbose:
+            print("Generating sequence of max length:", seq_len)
+    
+        # kv caches
+    
+        cache = None
+    
+        # if doing contrastive decoding, turn off filter automatically
+    
+        if exists(amateur_model):
+            amateur_model = cast_tuple(amateur_model)
+            contrastive_decode_kwargs = cast_tuple(contrastive_decode_kwargs)
+    
+            assert len(amateur_model) == len(contrastive_decode_kwargs)
+    
+            amateur_caches = [None] * len(amateur_model)
+            filter_logits_fn = identity
+    
+            for i, module in enumerate(amateur_model):
+                if isinstance(module, AutoregressiveWrapper):
+                    amateur_model[i] = module.net
+    
+                module.eval()
+    
+        # precompute some tensors for type controls
+    
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids.to(device)
+    
+            # build per-token temperature and bias vectors if provided
+            per_token_temp = None
+            if type_temperatures is not None and len(type_temperatures) > 0:
+                per_token_temp = torch.ones_like(token_type_ids, dtype=torch.float32)
+                for type_id, temp_val in type_temperatures.items():
+                    per_token_temp[token_type_ids == type_id] = float(temp_val)
+    
+            per_token_bias = None
+            if type_biases is not None and len(type_biases) > 0:
+                per_token_bias = torch.zeros_like(token_type_ids, dtype=torch.float32)
+                for type_id, bias_val in type_biases.items():
+                    per_token_bias[token_type_ids == type_id] = float(bias_val)
+    
+            # repetition penalty per type
+            per_type_rep_penalty = repetition_penalty_per_type or {}
+    
+            # rare type mask
+            rare_type_mask = None
+            if rare_types is not None and len(rare_types) > 0:
+                rare_type_mask = torch.zeros_like(token_type_ids, dtype=torch.bool)
+                for rt in rare_types:
+                    rare_type_mask |= (token_type_ids == rt)
+        else:
+            per_token_temp = None
+            per_token_bias = None
+            per_type_rep_penalty = {}
+            rare_type_mask = None
+    
+        # prepare forbidden mask if provided
+        # We'll lazily convert forbidden_token_ids into a boolean mask of shape [b, vocab]
+        forbidden_mask_per_batch = None
+        if forbidden_token_ids is not None:
+            # If it's a LongTensor of ids (1D)
+            if forbidden_token_ids.dtype in (torch.int64, torch.int32):
+                # create a [vocab] bool mask from ids
+                vocab_size = self.net.config.vocab_size if hasattr(self.net, 'config') else None
+                # If we can't infer vocab_size, we'll infer from token_type_ids if available
+                if vocab_size is None and token_type_ids is not None:
+                    vocab_size = token_type_ids.shape[0]
+                assert vocab_size is not None, "Cannot infer vocab size for forbidden_token_ids; provide a boolean mask instead."
+                mask = torch.zeros(vocab_size, dtype=torch.bool, device=device)
+                ids = forbidden_token_ids.to(device)
+                mask[ids.clamp(0, vocab_size-1)] = True
+                forbidden_mask_per_batch = mask.unsqueeze(0).expand(b, -1)  # [b, vocab]
+            elif forbidden_token_ids.dtype == torch.bool:
+                # could be [vocab] or [b, vocab]
+                if forbidden_token_ids.dim() == 1:
+                    forbidden_mask_per_batch = forbidden_token_ids.to(device).unsqueeze(0).expand(b, -1)
+                elif forbidden_token_ids.dim() == 2:
+                    assert forbidden_token_ids.shape[0] == b, "forbidden_token_ids batch dimension must match prompts batch size"
+                    forbidden_mask_per_batch = forbidden_token_ids.to(device)
+                else:
+                    raise ValueError("forbidden_token_ids boolean mask must be 1D [vocab] or 2D [b, vocab]")
+            else:
+                raise TypeError("forbidden_token_ids must be LongTensor of ids or BoolTensor mask")
+    
+        # sampling up to seq_len
+    
+        for sl in range(seq_len):
+    
+            if restrict_to_max_seq_len:
+                max_len_exceeded = out.shape[-1] > max_seq_len
+    
+                assert not (cache_kv and max_len_exceeded and not self.net.can_cache_kv_outside_max_seq_len), \
+                    'the network cannot use cached key values when decoding outside the max sequence length. ' \
+                    'most likely because you are using absolute positional embedding. ' \
+                    'you can switch to rotary embeddings to resolve this issue'
+    
+                x = out[:, -max_seq_len:]
+    
+                if exists(cache):
+                    for inter in cache.attn_intermediates:
+                        if inter.layer_type == 'a':
+                            inter.cached_kv = [t[..., -(max_seq_len - 1):, :] for t in inter.cached_kv]
+    
+            logits, new_cache = self.net(
+                x,
+                return_intermediates = True,
+                cache = cache,
+                seq_start_pos = seq_start_pos,
+                **kwargs
+            )
+    
+            if cache_kv and self.net.can_cache_kv:
+                cache = new_cache
+    
+            logits = logits[:, -1]  # [b, vocab]
+    
+            # handle contrastive decoding
+    
+            if exists(amateur_model):
+                for i, (amateur, amateur_cache, amateur_contrastive_decode_kwargs) in enumerate(
+                    zip(amateur_model, amateur_caches, contrastive_decode_kwargs)
+                ):
+                    amateur_logits, next_amateur_cache = amateur(
+                        x,
+                        return_intermediates = True,
+                        cache = amateur_cache,
+                        seq_start_pos = seq_start_pos,
+                        **kwargs
+                    )
+    
+                    amateur_logits = amateur_logits[:, -1]
+    
+                    assert amateur_logits.shape == logits.shape, \
+                        'logits dimension are not the same between amateur and expert model'
+                    logits = contrastive_decode_fn(logits, amateur_logits, **amateur_contrastive_decode_kwargs)
+    
+                    if cache_kv and amateur.can_cache_kv:
+                        amateur_caches[i] = next_amateur_cache
+    
+            # --------- STRUCTURED LOGIT SHAPING (no training) ---------
+    
+            if token_type_ids is not None:
+    
+                # 1) per-token bias (type-aware)
+                if per_token_bias is not None:
+                    logits = logits + per_token_bias  # broadcast [vocab]
+    
+                # 2) repetition penalty per type (context-aware)
+                if repetition_window > 0 and len(per_type_rep_penalty) > 0:
+                    # look at recent tokens
+                    recent = out[:, -repetition_window:].to(device)  # [b, w]
+                    # map to types
+                    recent_types = token_type_ids[recent]  # [b, w]
+    
+                    # for each type, compute frequency and apply penalty
+                    # we do this per batch element
+                    for bi in range(b):
+                        types_b = recent_types[bi]  # [w]
+                        if types_b.numel() == 0:
+                            continue
+                        # count occurrences per type id present in penalties
+                        for type_id, penalty_scale in per_type_rep_penalty.items():
+                            # penalty_scale > 1.0 means stronger penalty
+                            mask = (types_b == type_id)
+                            if mask.any():
+                                freq = mask.float().mean().item()  # 0..1
+                                if freq > 0.0:
+                                    # build a penalty vector for this type
+                                    type_mask = (token_type_ids == type_id)  # [vocab]
+                                    # subtract a penalty proportional to freq
+                                    # (log-space penalty)
+                                    logits[bi, type_mask] /= (1.0 + freq * (penalty_scale - 1.0))
+    
+                # 3) entropy-based rare-type boost (gentle, context-aware)
+                if rare_type_mask is not None and rare_type_boost > 0.0:
+                    # compute current probs & entropy (before global temperature)
+                    probs_raw = F.softmax(logits, dim=-1)  # [b, vocab]
+                    log_probs_raw = torch.log(probs_raw + 1e-9)
+                    entropy = -(probs_raw * log_probs_raw).sum(dim=-1)  # [b]
+    
+                    # for low-entropy states, gently boost rare types
+                    low_entropy = entropy < entropy_threshold
+                    if low_entropy.any():
+                        # boost only for those batch elements
+                        boost_vec = torch.zeros_like(logits)
+                        boost_vec[:, rare_type_mask] = rare_type_boost
+                        logits = torch.where(
+                            low_entropy.unsqueeze(-1),
+                            logits + boost_vec,
+                            logits
+                        )
+    
+                # 4) per-token temperature (type-aware)
+                # apply before global temperature
+                if per_token_temp is not None:
+                    # divide logits by per-token temperature
+                    # (smaller temp -> sharper distribution for that type)
+                    logits = logits / per_token_temp
+    
+            # --------- APPLY FORBIDDEN TOKEN MASK ---------
+            if forbidden_mask_per_batch is not None:
+                # ensure shapes match
+                assert forbidden_mask_per_batch.shape[0] == b and forbidden_mask_per_batch.shape[1] == logits.shape[-1], \
+                    "forbidden mask shape must be [b, vocab]"
+                # set logits for forbidden tokens to a large negative value
+                logits = logits.masked_fill(forbidden_mask_per_batch, float(forbidden_value))
+    
+            # ----------------------------------------------------------
+    
+            # filter by top_k, top_p (nucleus), top_a, or custom
+    
+            if greedy:
+                sample = logits.argmax(dim = -1, keepdim = True)
+            else:
+                filtered_logits = filter_logits_fn(logits, **filter_kwargs)
+                probs = F.softmax(filtered_logits / temperature, dim=-1)
+                sample = torch.multinomial(probs, 1)
+    
+            # concat sample
+    
+            out = torch.cat((out, sample), dim=-1)
+    
+            if verbose:
+                if sl % 32 == 0:
+                    print(sl, '/', seq_len)
+    
+            if not exists(eos_token):
+                continue
+    
+            is_eos_tokens = (out == eos_token)
+    
+            if is_eos_tokens.any(dim = -1).all():
+    
+                if verbose:
+                    print('Model called the end of sequence at:', sl, '/', seq_len)
+    
+                break
+    
+        if exists(eos_token):
+            # mask out everything after the eos tokens
+            shifted_is_eos_tokens = F.pad(is_eos_tokens, (1, -1))
+            mask = shifted_is_eos_tokens.float().cumsum(dim = -1) >= 1
+            out = out.masked_fill(mask, self.pad_value)
+    
+        if return_prime:
+            out = out[:, :]
+        else:
+            out = out[:, t:]
+    
+        out, = unpack(out, ps, '* n')
+    
+        return out
+
+#=========================================================================================
+
+# Binary classifier fuctions
+
+class ClsInferenceDataset(Dataset):
+    """
+    Dataset for pairs (src_seq, label).
+    src_seq: list of token IDs (ints).
+    label: single int or float (0 or 1).
+    """
+    def __init__(self, data_pairs):
+        self.data_pairs = data_pairs
+
+    def __len__(self):
+        return len(self.data_pairs)
+
+    def __getitem__(self, idx):
+        src_seq = self.data_pairs[idx]
+        x = torch.tensor(src_seq, dtype=torch.long)
+        return x
+
+def build_cls_model(num_tokens=18819,
+                    max_seq_len=1024,
+                    logits_dim=1,
+                    use_cls_token=True,
+                    squeeze_out_last_dim=True,
+                    dim=1024,
+                    depth=8,
+                    heads=8,
+                    device='cuda'
+                   ):
+
+    """
+    Constructs the Transformer model that outputs a single logit per input.
+    """
+
+    model = TransformerWrapper(
+        num_tokens=num_tokens,
+        max_seq_len=max_seq_len,
+        logits_dim=logits_dim,
+        use_cls_token=use_cls_token,
+        squeeze_out_last_dim = squeeze_out_last_dim,
+        attn_layers=Encoder(dim=dim,
+                            depth=depth,
+                            heads=heads
+                           )
+    )
+
+    return model.to(device)
+
+def load_cls_model(checkpoint_path, device='cuda'):
+    
+    """
+    Rebuilds the architecture, loads weights.
+    """
+    
+    model = build_cls_model(device=device)
+    state = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(state)
+    model.to(device).eval()
+    
+    return model
+
+def cls_predict(model,
+                seqs,
+                batch_size=8,
+                threshold=0.5,
+                seq_len=1024,
+                pad_token=18818,
+                device='cuda'
+               ):
+    
+    """
+    Returns two lists:
+      - probs: float probabilities  
+      - preds: int 0/1 predictions  
+    """
+    
+    def collate_fn(batch):
+        # batch: list of sequences (list/1D-tensor)
+        tensors = [s[:seq_len].detach().clone() for s in batch]
+        max_len = min(seq_len, max(t.size(0) for t in tensors))
+        padded = torch.full((len(tensors), max_len), pad_token, dtype=torch.long)
+        for i, t in enumerate(tensors):
+            L = t.size(0)
+            padded[i, :L] = t
+        return padded
+
+    ds = ClsInferenceDataset(seqs)
+    loader = DataLoader(ds, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
+
+    all_probs = []
+    all_preds = []
+
+    model.to(device)
+    model.eval()
+    
+    with torch.inference_mode():
+        for x in loader:
+            
+            x = x.to(device)                       # [B, L] (truncated & padded)
+            
+            logits = model(x).squeeze()            # [B]
+            
+            probs = torch.sigmoid(logits)         # [B]
+            
+            preds = (probs >= threshold).long()
+
+            probs = probs.cpu().tolist()
+            preds = preds.cpu().tolist()
+
+            if type(preds) == list:
+                all_probs.extend(probs)
+                all_preds.extend(preds)
+
+            else:
+                all_probs.append(probs)
+                all_preds.append(preds)                
+
+    return all_preds, all_probs
+
 #=================================================================================================================================
 # This is the end of x_transformer_2_3_1 Python module
 #=================================================================================================================================