class_docstr: Transformer-based language model on multiple streams of codes.
  
  Args:
      pattern_provider (CodebooksPatternProvider): Pattern provider for codebook interleaving.
      condition_provider (MusicConditioningProvider): Conditioning provider from metadata.
      fuser (ConditionFuser): Fuser handling the fusing of conditions with language model input.
      n_q (int): Number of parallel streams to model.
      card (int): Cardinality, vocabulary size.
      dim (int): Dimension of the transformer encoder.
      num_heads (int): Number of heads for the transformer encoder.
      hidden_scale (int): Scale for hidden feed forward dimension of the transformer encoder.
      norm (str): Normalization method.
      norm_first (bool): Use pre-norm instead of post-norm.
      emb_lr (float, optional): Embedding-specific learning rate.
      bias_proj (bool): Use bias for output projections.
      weight_init (str, optional): Method for weight initialization.
      depthwise_init (str, optional): Method for depthwise weight initialization.
      zero_bias_init (bool): If true and bias in Linears, initialize bias to zeros.
      cfg_dropout (float): Classifier-free guidance dropout.
      cfg_coef (float): Classifier-free guidance coefficient.
      attribute_dropout (dict): Attribute dropout probabilities.
      two_step_cfg (bool): Whether to run classifier free-guidance with 2 distinct steps.
      **kwargs: Additional parameters for the transformer encoder.
docstr: Sample next token from the model given a sequence and a set of conditions. The model supports
  multiple sampling strategies (greedy sampling, softmax, top-k, top-p...).
  
  Args:
      sequence (torch.Tensor): Current sequence of shape [B, K, S]
          with K corresponding to the number of codebooks and S the number of sequence steps.
          S = 1 in streaming mode, except for the first step that contains a bigger prompt.
      condition_tensors (dict[str, ConditionType): Set of conditions. If CFG is used,
          should be twice the batch size, being the concatenation of the conditions + null conditions.
      use_sampling (bool): Whether to use a sampling strategy or not.
      temp (float): Sampling temperature.
      top_k (int): K for "top-k" sampling.
      top_p (float): P for "top-p" sampling.
      cfg_coef (float, optional): classifier free guidance coefficient
  Returns:
      next_token (torch.Tensor): Next token tensor of shape [B, K, 1].
docstr_tok: ["Sample","next","token","from","the","model","given","a","sequence","and","a","set","of","conditions",".","The","model","supports","multiple","sampling","strategies","(","greedy","sampling",",","softmax",",","top-k",",","top-p","...",")",".","Args",":","sequence","(","torch.Tensor",")",":","Current","sequence","of","shape","[","B",",","K",",","S","]","with","K","corresponding","to","the","number","of","codebooks","and","S","the","number","of","sequence","steps",".","S","=","1","in","streaming","mode",",","except","for","the","first","step","that","contains","a","bigger","prompt",".","condition_tensors","(","dict","[","str",",","ConditionType",")",":","Set","of","conditions",".","If","CFG","is","used",",","should","be","twice","the","batch","size",",","being","the","concatenation","of","the","conditions","+","null","conditions",".","use_sampling","(","bool",")",":","Whether","to","use","a","sampling","strategy","or","not",".","temp","(","float",")",":","Sampling","temperature",".","top_k","(","int",")",":","K","for","``","top-k","''","sampling",".","top_p","(","float",")",":","P","for","``","top-p","''","sampling",".","cfg_coef","(","float",",","optional",")",":","classifier","free","guidance","coefficient","Returns",":","next_token","(","torch.Tensor",")",":","Next","token","tensor","of","shape","[","B",",","K",",","1","]","."]
code: def _sample_next_token(
      self,
      sequence: torch.Tensor,
      cfg_conditions: CFGConditions,
      unconditional_state: State,
      use_sampling: bool = False,
      temp: float = 1.0,
      top_k: int = 0,
      top_p: float = 0.0,
      cfg_coef: tp.Optional[float] = None,
      two_step_cfg: tp.Optional[bool] = None,
  ) -> torch.Tensor:
      """Sample next token from the model given a sequence and a set of conditions. The model supports
      multiple sampling strategies (greedy sampling, softmax, top-k, top-p...).
  
      Args:
          sequence (torch.Tensor): Current sequence of shape [B, K, S]
              with K corresponding to the number of codebooks and S the number of sequence steps.
              S = 1 in streaming mode, except for the first step that contains a bigger prompt.
          condition_tensors (dict[str, ConditionType): Set of conditions. If CFG is used,
              should be twice the batch size, being the concatenation of the conditions + null conditions.
          use_sampling (bool): Whether to use a sampling strategy or not.
          temp (float): Sampling temperature.
          top_k (int): K for "top-k" sampling.
          top_p (float): P for "top-p" sampling.
          cfg_coef (float, optional): classifier free guidance coefficient
      Returns:
          next_token (torch.Tensor): Next token tensor of shape [B, K, 1].
      """
      B = sequence.shape[0]
      cfg_coef = self.cfg_coef if cfg_coef is None else cfg_coef
      model = self if self._fsdp is None else self._fsdp
      two_step_cfg = (
          self.two_step_cfg if two_step_cfg is None else two_step_cfg
      )
      if two_step_cfg and cfg_conditions != {}:
          assert isinstance(cfg_conditions, tuple), type(cfg_conditions)
          condition_tensors, null_condition_tensors = cfg_conditions
          cond_logits = model(
              sequence,
              conditions=[],
              condition_tensors=condition_tensors,
          )
          state = self.get_streaming_state()
          self.set_streaming_state(unconditional_state)
          uncond_logits = model(
              sequence,
              conditions=[],
              condition_tensors=null_condition_tensors,
          )
          unconditional_state.update(self.get_streaming_state())
          self.set_streaming_state(state)
          logits = (
              uncond_logits
              + (cond_logits - uncond_logits) * self.cfg_coef
          )
      else:
          assert isinstance(cfg_conditions, dict)
          condition_tensors = cfg_conditions
          if condition_tensors:
              # Preparing for CFG, predicting both conditional and unconditional logits.
              sequence = torch.cat([sequence, sequence], dim=0)
          all_logits = model(
              sequence,
              conditions=[],
              condition_tensors=condition_tensors,
          )
          if condition_tensors:
              cond_logits, uncond_logits = all_logits.split(
                  B, dim=0
              )  # [B, K, T, card]
              logits = (
                  uncond_logits
                  + (cond_logits - uncond_logits) * cfg_coef
              )
          else:
              logits = all_logits
  
      logits = logits.permute(0, 1, 3, 2)  # [B, K, card, T]
      logits = logits[..., -1]  # [B x K x card]
  
      # Apply softmax for sampling if temp > 0. Else, do greedy sampling to avoid zero division error.
      if use_sampling and temp > 0.0:
          probs = torch.softmax(logits / temp, dim=-1)
          if top_p > 0.0:
              next_token = utils.sample_top_p(probs, p=top_p)
          elif top_k > 0:
              next_token = utils.sample_top_k(probs, k=top_k)
          else:
              next_token = utils.multinomial(probs, num_samples=1)
      else:
          next_token = torch.argmax(logits, dim=-1, keepdim=True)
  
      return next_token
  
imports: from dataclasses import dataclass
  from functools import partial
  import logging
  import math
  import typing
  import torch
  from torch import nn
  from ..utils import utils
  from ..modules.streaming import StreamingModule, State
  from ..modules.transformer import StreamingTransformer, create_norm_fn
  from ..modules.conditioners import ConditionFuser, ClassifierFreeGuidanceDropout, AttributeDropout, ConditioningProvider, ConditioningAttributes, ConditionType
  from ..modules.codebooks_patterns import CodebooksPatternProvider
  from ..modules.activations import get_activation_fn
label_desc: Use image node_id 7 for calling the LMModel obj's underlying member method code with example usage: obj._sample_next_token(sequence, cfg_conditions, unconditional_state, use_sampling, temp, top_k, top_p, cfg_coef, two_step_cfg) and returns: next_token