spark-nlp 5.5.3__py2.py3-none-any.whl → 6.0.1rc1__py2.py3-none-any.whl

sparknlp/common/properties.py

@@ -12,6 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 """Contains classes for Annotator properties."""
+from typing import List, Dict
 
 from pyspark.ml.param import Param, Params, TypeConverters
 
@@ -628,133 +629,641 @@ class HasGeneratorProperties:
                              typeConverter=TypeConverters.toInt)
 
 
-def setTask(self, value):
-    """Sets the transformer's task, e.g. ``summarize:``.
+    def setTask(self, value):
+        """Sets the transformer's task, e.g. ``summarize:``.
 
-    Parameters
-    ----------
-    value : str
-        The transformer's task
-    """
-    return self._set(task=value)
+        Parameters
+        ----------
+        value : str
+            The transformer's task
+        """
+        return self._set(task=value)
 
 
-def setMinOutputLength(self, value):
-    """Sets minimum length of the sequence to be generated.
+    def setMinOutputLength(self, value):
+        """Sets minimum length of the sequence to be generated.
 
-    Parameters
-    ----------
-    value : int
-        Minimum length of the sequence to be generated
-    """
-    return self._set(minOutputLength=value)
+        Parameters
+        ----------
+        value : int
+            Minimum length of the sequence to be generated
+        """
+        return self._set(minOutputLength=value)
 
 
-def setMaxOutputLength(self, value):
-    """Sets maximum length of output text.
+    def setMaxOutputLength(self, value):
+        """Sets maximum length of output text.
 
-    Parameters
-    ----------
-    value : int
-        Maximum length of output text
-    """
-    return self._set(maxOutputLength=value)
+        Parameters
+        ----------
+        value : int
+            Maximum length of output text
+        """
+        return self._set(maxOutputLength=value)
 
 
-def setDoSample(self, value):
-    """Sets whether or not to use sampling, use greedy decoding otherwise.
+    def setDoSample(self, value):
+        """Sets whether or not to use sampling, use greedy decoding otherwise.
 
-    Parameters
-    ----------
-    value : bool
-        Whether or not to use sampling; use greedy decoding otherwise
-    """
-    return self._set(doSample=value)
+        Parameters
+        ----------
+        value : bool
+            Whether or not to use sampling; use greedy decoding otherwise
+        """
+        return self._set(doSample=value)
 
 
-def setTemperature(self, value):
-    """Sets the value used to module the next token probabilities.
+    def setTemperature(self, value):
+        """Sets the value used to module the next token probabilities.
 
-    Parameters
-    ----------
-    value : float
-        The value used to module the next token probabilities
-    """
-    return self._set(temperature=value)
+        Parameters
+        ----------
+        value : float
+            The value used to module the next token probabilities
+        """
+        return self._set(temperature=value)
 
 
-def setTopK(self, value):
-    """Sets the number of highest probability vocabulary tokens to keep for
-    top-k-filtering.
+    def setTopK(self, value):
+        """Sets the number of highest probability vocabulary tokens to keep for
+        top-k-filtering.
 
-    Parameters
-    ----------
-    value : int
-        Number of highest probability vocabulary tokens to keep
-    """
-    return self._set(topK=value)
+        Parameters
+        ----------
+        value : int
+            Number of highest probability vocabulary tokens to keep
+        """
+        return self._set(topK=value)
 
 
-def setTopP(self, value):
-    """Sets the top cumulative probability for vocabulary tokens.
+    def setTopP(self, value):
+        """Sets the top cumulative probability for vocabulary tokens.
 
-    If set to float < 1, only the most probable tokens with probabilities
-    that add up to ``topP`` or higher are kept for generation.
+        If set to float < 1, only the most probable tokens with probabilities
+        that add up to ``topP`` or higher are kept for generation.
 
-    Parameters
-    ----------
-    value : float
-        Cumulative probability for vocabulary tokens
-    """
-    return self._set(topP=value)
+        Parameters
+        ----------
+        value : float
+            Cumulative probability for vocabulary tokens
+        """
+        return self._set(topP=value)
 
 
-def setRepetitionPenalty(self, value):
-    """Sets the parameter for repetition penalty. 1.0 means no penalty.
+    def setRepetitionPenalty(self, value):
+        """Sets the parameter for repetition penalty. 1.0 means no penalty.
 
-    Parameters
-    ----------
-    value : float
-        The repetition penalty
+        Parameters
+        ----------
+        value : float
+            The repetition penalty
 
-    References
-    ----------
-    See `Ctrl: A Conditional Transformer Language Model For Controllable
-    Generation <https://arxiv.org/pdf/1909.05858.pdf>`__ for more details.
-    """
-    return self._set(repetitionPenalty=value)
+        References
+        ----------
+        See `Ctrl: A Conditional Transformer Language Model For Controllable
+        Generation <https://arxiv.org/pdf/1909.05858.pdf>`__ for more details.
+        """
+        return self._set(repetitionPenalty=value)
 
 
-def setNoRepeatNgramSize(self, value):
-    """Sets size of n-grams that can only occur once.
+    def setNoRepeatNgramSize(self, value):
+        """Sets size of n-grams that can only occur once.
 
-    If set to int > 0, all ngrams of that size can only occur once.
+        If set to int > 0, all ngrams of that size can only occur once.
 
-    Parameters
-    ----------
-    value : int
-        N-gram size can only occur once
-    """
-    return self._set(noRepeatNgramSize=value)
+        Parameters
+        ----------
+        value : int
+            N-gram size can only occur once
+        """
+        return self._set(noRepeatNgramSize=value)
+
+
+    def setBeamSize(self, value):
+        """Sets the number of beam size for beam search.
+
+        Parameters
+        ----------
+        value : int
+            Number of beam size for beam search
+        """
+        return self._set(beamSize=value)
+
+
+    def setNReturnSequences(self, value):
+        """Sets the number of sequences to return from the beam search.
+
+        Parameters
+        ----------
+        value : int
+            Number of sequences to return
+        """
+        return self._set(nReturnSequences=value)
+
+
+class HasLlamaCppProperties:
+    # -------- MODEl PARAMETERS --------
+    nThreads = Param(Params._dummy(), "nThreads", "Set the number of threads to use during generation",
+                     typeConverter=TypeConverters.toInt)
+    nThreadsDraft = Param(Params._dummy(), "nThreadsDraft", "Set the number of threads to use during draft generation",
+                          typeConverter=TypeConverters.toInt)
+    nThreadsBatch = Param(Params._dummy(), "nThreadsBatch",
+                          "Set the number of threads to use during batch and prompt processing",
+                          typeConverter=TypeConverters.toInt)
+    nThreadsBatchDraft = Param(Params._dummy(), "nThreadsBatchDraft",
+                               "Set the number of threads to use during batch and prompt processing",
+                               typeConverter=TypeConverters.toInt)
+    nCtx = Param(Params._dummy(), "nCtx", "Set the size of the prompt context", typeConverter=TypeConverters.toInt)
+    nBatch = Param(Params._dummy(), "nBatch",
+                   "Set the logical batch size for prompt processing (must be >=32 to use BLAS)",
+                   typeConverter=TypeConverters.toInt)
+    nUbatch = Param(Params._dummy(), "nUbatch",
+                    "Set the physical batch size for prompt processing (must be >=32 to use BLAS)",
+                    typeConverter=TypeConverters.toInt)
+    nDraft = Param(Params._dummy(), "nDraft", "Set the number of tokens to draft for speculative decoding",
+                   typeConverter=TypeConverters.toInt)
+    nChunks = Param(Params._dummy(), "nChunks", "Set the maximal number of chunks to process",
+                    typeConverter=TypeConverters.toInt)
+    nSequences = Param(Params._dummy(), "nSequences", "Set the number of sequences to decode",
+                       typeConverter=TypeConverters.toInt)
+    pSplit = Param(Params._dummy(), "pSplit", "Set the speculative decoding split probability",
+                   typeConverter=TypeConverters.toFloat)
+    nGpuLayers = Param(Params._dummy(), "nGpuLayers", "Set the number of layers to store in VRAM (-1 - use default)",
+                       typeConverter=TypeConverters.toInt)
+    nGpuLayersDraft = Param(Params._dummy(), "nGpuLayersDraft",
+                            "Set the number of layers to store in VRAM for the draft model (-1 - use default)",
+                            typeConverter=TypeConverters.toInt)
+    # Set how to split the model across GPUs
+    #
+    #   - NONE: No GPU split
+    #   - LAYER: Split the model across GPUs by layer
+    #   - ROW: Split the model across GPUs by rows
+    gpuSplitMode = Param(Params._dummy(), "gpuSplitMode", "Set how to split the model across GPUs",
+                         typeConverter=TypeConverters.toString)
+    mainGpu = Param(Params._dummy(), "mainGpu", "Set the main GPU that is used for scratch and small tensors.",
+                    typeConverter=TypeConverters.toInt)
+    tensorSplit = Param(Params._dummy(), "tensorSplit", "Set how split tensors should be distributed across GPUs",
+                        typeConverter=TypeConverters.toListFloat)
+    grpAttnN = Param(Params._dummy(), "grpAttnN", "Set the group-attention factor", typeConverter=TypeConverters.toInt)
+    grpAttnW = Param(Params._dummy(), "grpAttnW", "Set the group-attention width", typeConverter=TypeConverters.toInt)
+    ropeFreqBase = Param(Params._dummy(), "ropeFreqBase", "Set the RoPE base frequency, used by NTK-aware scaling",
+                         typeConverter=TypeConverters.toFloat)
+    ropeFreqScale = Param(Params._dummy(), "ropeFreqScale",
+                          "Set the RoPE frequency scaling factor, expands context by a factor of 1/N",
+                          typeConverter=TypeConverters.toFloat)
+    yarnExtFactor = Param(Params._dummy(), "yarnExtFactor", "Set the YaRN extrapolation mix factor",
+                          typeConverter=TypeConverters.toFloat)
+    yarnAttnFactor = Param(Params._dummy(), "yarnAttnFactor", "Set the YaRN scale sqrt(t) or attention magnitude",
+                           typeConverter=TypeConverters.toFloat)
+    yarnBetaFast = Param(Params._dummy(), "yarnBetaFast", "Set the YaRN low correction dim or beta",
+                         typeConverter=TypeConverters.toFloat)
+    yarnBetaSlow = Param(Params._dummy(), "yarnBetaSlow", "Set the YaRN high correction dim or alpha",
+                         typeConverter=TypeConverters.toFloat)
+    yarnOrigCtx = Param(Params._dummy(), "yarnOrigCtx", "Set the YaRN original context size of model",
+                        typeConverter=TypeConverters.toInt)
+    defragmentationThreshold = Param(Params._dummy(), "defragmentationThreshold",
+                                     "Set the KV cache defragmentation threshold", typeConverter=TypeConverters.toFloat)
+    # Set optimization strategies that help on some NUMA systems (if available)
+    #
+    # Available Strategies:
+    #
+    #   - DISABLED: No NUMA optimizations
+    #   - DISTRIBUTE: Spread execution evenly over all
+    #   - ISOLATE: Only spawn threads on CPUs on the node that execution started on
+    #   - NUMA_CTL: Use the CPU map provided by numactl
+    #   - MIRROR: Mirrors the model across NUMA nodes
+    numaStrategy = Param(Params._dummy(), "numaStrategy",
+                         "Set optimization strategies that help on some NUMA systems (if available)",
+                         typeConverter=TypeConverters.toString)
+    # Set the RoPE frequency scaling method, defaults to linear unless specified by the model.
+    #
+    #   - UNSPECIFIED: Don't use any scaling
+    #   - LINEAR: Linear scaling
+    #   - YARN: YaRN RoPE scaling
+    ropeScalingType = Param(Params._dummy(), "ropeScalingType",
+                            "Set the RoPE frequency scaling method, defaults to linear unless specified by the model",
+                            typeConverter=TypeConverters.toString)
+    # Set the pooling type for embeddings, use model default if unspecified
+    #
+    #   - 0 NONE: Don't use any pooling
+    #   - 1 MEAN: Mean Pooling
+    #   - 2 CLS: CLS Pooling
+    poolingType = Param(Params._dummy(), "poolingType",
+                        "Set the pooling type for embeddings, use model default if unspecified",
+                        typeConverter=TypeConverters.toString)
+    modelDraft = Param(Params._dummy(), "modelDraft", "Set the draft model for speculative decoding",
+                       typeConverter=TypeConverters.toString)
+    modelAlias = Param(Params._dummy(), "modelAlias", "Set a model alias", typeConverter=TypeConverters.toString)
+    lookupCacheStaticFilePath = Param(Params._dummy(), "lookupCacheStaticFilePath",
+                                      "Set path to static lookup cache to use for lookup decoding (not updated by generation)",
+                                      typeConverter=TypeConverters.toString)
+    lookupCacheDynamicFilePath = Param(Params._dummy(), "lookupCacheDynamicFilePath",
+                                       "Set path to dynamic lookup cache to use for lookup decoding (updated by generation)",
+                                       typeConverter=TypeConverters.toString)
+    # loraAdapters = new StructFeature[Map[String, Float]](this, "loraAdapters")
+    embedding = Param(Params._dummy(), "embedding", "Whether to load model with embedding support",
+                      typeConverter=TypeConverters.toBoolean)
+    flashAttention = Param(Params._dummy(), "flashAttention", "Whether to enable Flash Attention",
+                           typeConverter=TypeConverters.toBoolean)
+    inputPrefixBos = Param(Params._dummy(), "inputPrefixBos",
+                           "Whether to add prefix BOS to user inputs, preceding the `--in-prefix` string",
+                           typeConverter=TypeConverters.toBoolean)
+    useMmap = Param(Params._dummy(), "useMmap",
+                    "Whether to use memory-map model (faster load but may increase pageouts if not using mlock)",
+                    typeConverter=TypeConverters.toBoolean)
+    useMlock = Param(Params._dummy(), "useMlock",
+                     "Whether to force the system to keep model in RAM rather than swapping or compressing",
+                     typeConverter=TypeConverters.toBoolean)
+    noKvOffload = Param(Params._dummy(), "noKvOffload", "Whether to disable KV offload",
+                        typeConverter=TypeConverters.toBoolean)
+    systemPrompt = Param(Params._dummy(), "systemPrompt", "Set a system prompt to use",
+                         typeConverter=TypeConverters.toString)
+    chatTemplate = Param(Params._dummy(), "chatTemplate", "The chat template to use",
+                         typeConverter=TypeConverters.toString)
+
+    # -------- INFERENCE PARAMETERS --------
+    inputPrefix = Param(Params._dummy(), "inputPrefix", "Set the prompt to start generation with",
+                        typeConverter=TypeConverters.toString)
+    inputSuffix = Param(Params._dummy(), "inputSuffix", "Set a suffix for infilling",
+                        typeConverter=TypeConverters.toString)
+    cachePrompt = Param(Params._dummy(), "cachePrompt", "Whether to remember the prompt to avoid reprocessing it",
+                        typeConverter=TypeConverters.toBoolean)
+    nPredict = Param(Params._dummy(), "nPredict", "Set the number of tokens to predict",
+                     typeConverter=TypeConverters.toInt)
+    topK = Param(Params._dummy(), "topK", "Set top-k sampling", typeConverter=TypeConverters.toInt)
+    topP = Param(Params._dummy(), "topP", "Set top-p sampling", typeConverter=TypeConverters.toFloat)
+    minP = Param(Params._dummy(), "minP", "Set min-p sampling", typeConverter=TypeConverters.toFloat)
+    tfsZ = Param(Params._dummy(), "tfsZ", "Set tail free sampling, parameter z", typeConverter=TypeConverters.toFloat)
+    typicalP = Param(Params._dummy(), "typicalP", "Set locally typical sampling, parameter p",
+                     typeConverter=TypeConverters.toFloat)
+    temperature = Param(Params._dummy(), "temperature", "Set the temperature", typeConverter=TypeConverters.toFloat)
+    dynamicTemperatureRange = Param(Params._dummy(), "dynatempRange", "Set the dynamic temperature range",
+                                    typeConverter=TypeConverters.toFloat)
+    dynamicTemperatureExponent = Param(Params._dummy(), "dynatempExponent", "Set the dynamic temperature exponent",
+                                       typeConverter=TypeConverters.toFloat)
+    repeatLastN = Param(Params._dummy(), "repeatLastN", "Set the last n tokens to consider for penalties",
+                        typeConverter=TypeConverters.toInt)
+    repeatPenalty = Param(Params._dummy(), "repeatPenalty", "Set the penalty of repeated sequences of tokens",
+                          typeConverter=TypeConverters.toFloat)
+    frequencyPenalty = Param(Params._dummy(), "frequencyPenalty", "Set the repetition alpha frequency penalty",
+                             typeConverter=TypeConverters.toFloat)
+    presencePenalty = Param(Params._dummy(), "presencePenalty", "Set the repetition alpha presence penalty",
+                            typeConverter=TypeConverters.toFloat)
+    miroStat = Param(Params._dummy(), "miroStat", "Set MiroStat sampling strategies.",
+                     typeConverter=TypeConverters.toString)
+    miroStatTau = Param(Params._dummy(), "mirostatTau", "Set the MiroStat target entropy, parameter tau",
+                        typeConverter=TypeConverters.toFloat)
+    miroStatEta = Param(Params._dummy(), "mirostatEta", "Set the MiroStat learning rate, parameter eta",
+                        typeConverter=TypeConverters.toFloat)
+    penalizeNl = Param(Params._dummy(), "penalizeNl", "Whether to penalize newline tokens",
+                       typeConverter=TypeConverters.toBoolean)
+    nKeep = Param(Params._dummy(), "nKeep", "Set the number of tokens to keep from the initial prompt",
+                  typeConverter=TypeConverters.toInt)
+    seed = Param(Params._dummy(), "seed", "Set the RNG seed", typeConverter=TypeConverters.toInt)
+    nProbs = Param(Params._dummy(), "nProbs", "Set the amount top tokens probabilities to output if greater than 0.",
+                   typeConverter=TypeConverters.toInt)
+    minKeep = Param(Params._dummy(), "minKeep",
+                    "Set the amount of tokens the samplers should return at least (0 = disabled)",
+                    typeConverter=TypeConverters.toInt)
+    grammar = Param(Params._dummy(), "grammar", "Set BNF-like grammar to constrain generations",
+                    typeConverter=TypeConverters.toString)
+    penaltyPrompt = Param(Params._dummy(), "penaltyPrompt",
+                          "Override which part of the prompt is penalized for repetition.",
+                          typeConverter=TypeConverters.toString)
+    ignoreEos = Param(Params._dummy(), "ignoreEos",
+                      "Set whether to ignore end of stream token and continue generating (implies --logit-bias 2-inf)",
+                      typeConverter=TypeConverters.toBoolean)
+    disableTokenIds = Param(Params._dummy(), "disableTokenIds", "Set the token ids to disable in the completion",
+                            typeConverter=TypeConverters.toListInt)
+    stopStrings = Param(Params._dummy(), "stopStrings", "Set strings upon seeing which token generation is stopped",
+                        typeConverter=TypeConverters.toListString)
+    samplers = Param(Params._dummy(), "samplers", "Set which samplers to use for token generation in the given order",
+                     typeConverter=TypeConverters.toListString)
+    useChatTemplate = Param(Params._dummy(), "useChatTemplate",
+                            "Set whether or not generate should apply a chat template",
+                            typeConverter=TypeConverters.toBoolean)
+
+    # -------- MODEL SETTERS --------
+    def setNThreads(self, nThreads: int):
+        """Set the number of threads to use during generation"""
+        return self._set(nThreads=nThreads)
+
+    def setNThreadsDraft(self, nThreadsDraft: int):
+        """Set the number of threads to use during draft generation"""
+        return self._set(nThreadsDraft=nThreadsDraft)
+
+    def setNThreadsBatch(self, nThreadsBatch: int):
+        """Set the number of threads to use during batch and prompt processing"""
+        return self._set(nThreadsBatch=nThreadsBatch)
+
+    def setNThreadsBatchDraft(self, nThreadsBatchDraft: int):
+        """Set the number of threads to use during batch and prompt processing"""
+        return self._set(nThreadsBatchDraft=nThreadsBatchDraft)
+
+    def setNCtx(self, nCtx: int):
+        """Set the size of the prompt context"""
+        return self._set(nCtx=nCtx)
+
+    def setNBatch(self, nBatch: int):
+        """Set the logical batch size for prompt processing (must be >=32 to use BLAS)"""
+        return self._set(nBatch=nBatch)
+
+    def setNUbatch(self, nUbatch: int):
+        """Set the physical batch size for prompt processing (must be >=32 to use BLAS)"""
+        return self._set(nUbatch=nUbatch)
+
+    def setNDraft(self, nDraft: int):
+        """Set the number of tokens to draft for speculative decoding"""
+        return self._set(nDraft=nDraft)
+
+    def setNChunks(self, nChunks: int):
+        """Set the maximal number of chunks to process"""
+        return self._set(nChunks=nChunks)
+
+    def setNSequences(self, nSequences: int):
+        """Set the number of sequences to decode"""
+        return self._set(nSequences=nSequences)
+
+    def setPSplit(self, pSplit: float):
+        """Set the speculative decoding split probability"""
+        return self._set(pSplit=pSplit)
+
+    def setNGpuLayers(self, nGpuLayers: int):
+        """Set the number of layers to store in VRAM (-1 - use default)"""
+        return self._set(nGpuLayers=nGpuLayers)
+
+    def setNGpuLayersDraft(self, nGpuLayersDraft: int):
+        """Set the number of layers to store in VRAM for the draft model (-1 - use default)"""
+        return self._set(nGpuLayersDraft=nGpuLayersDraft)
+
+    def setGpuSplitMode(self, gpuSplitMode: str):
+        """Set how to split the model across GPUs"""
+        return self._set(gpuSplitMode=gpuSplitMode)
+
+    def setMainGpu(self, mainGpu: int):
+        """Set the main GPU that is used for scratch and small tensors."""
+        return self._set(mainGpu=mainGpu)
+
+    def setTensorSplit(self, tensorSplit: List[float]):
+        """Set how split tensors should be distributed across GPUs"""
+        return self._set(tensorSplit=tensorSplit)
+
+    def setGrpAttnN(self, grpAttnN: int):
+        """Set the group-attention factor"""
+        return self._set(grpAttnN=grpAttnN)
+
+    def setGrpAttnW(self, grpAttnW: int):
+        """Set the group-attention width"""
+        return self._set(grpAttnW=grpAttnW)
+
+    def setRopeFreqBase(self, ropeFreqBase: float):
+        """Set the RoPE base frequency, used by NTK-aware scaling"""
+        return self._set(ropeFreqBase=ropeFreqBase)
+
+    def setRopeFreqScale(self, ropeFreqScale: float):
+        """Set the RoPE frequency scaling factor, expands context by a factor of 1/N"""
+        return self._set(ropeFreqScale=ropeFreqScale)
+
+    def setYarnExtFactor(self, yarnExtFactor: float):
+        """Set the YaRN extrapolation mix factor"""
+        return self._set(yarnExtFactor=yarnExtFactor)
+
+    def setYarnAttnFactor(self, yarnAttnFactor: float):
+        """Set the YaRN scale sqrt(t) or attention magnitude"""
+        return self._set(yarnAttnFactor=yarnAttnFactor)
+
+    def setYarnBetaFast(self, yarnBetaFast: float):
+        """Set the YaRN low correction dim or beta"""
+        return self._set(yarnBetaFast=yarnBetaFast)
+
+    def setYarnBetaSlow(self, yarnBetaSlow: float):
+        """Set the YaRN high correction dim or alpha"""
+        return self._set(yarnBetaSlow=yarnBetaSlow)
+
+    def setYarnOrigCtx(self, yarnOrigCtx: int):
+        """Set the YaRN original context size of model"""
+        return self._set(yarnOrigCtx=yarnOrigCtx)
+
+    def setDefragmentationThreshold(self, defragmentationThreshold: float):
+        """Set the KV cache defragmentation threshold"""
+        return self._set(defragmentationThreshold=defragmentationThreshold)
 
+    def setNumaStrategy(self, numaStrategy: str):
+        """Set optimization strategies that help on some NUMA systems (if available)"""
+        numaUpper = numaStrategy.upper()
+        numaStrategies = ["DISABLED", "DISTRIBUTE", "ISOLATE", "NUMA_CTL", "MIRROR"]
+        if numaUpper not in numaStrategies:
+            raise ValueError(
+                f"Invalid NUMA strategy: {numaUpper}. "
+                + f"Valid values are: {numaStrategies}"
+            )
+        return self._set(numaStrategy=numaStrategy)
+
+    def setRopeScalingType(self, ropeScalingType: str):
+        """Set the RoPE frequency scaling method, defaults to linear unless specified by the model"""
+        return self._set(ropeScalingType=ropeScalingType)
+
+    def setPoolingType(self, poolingType: str):
+        """Set the pooling type for embeddings, use model default if unspecified"""
+        poolingTypeUpper = poolingType.upper()
+        poolingTypes = ["NONE", "MEAN", "CLS", "LAST"]
+        if poolingTypeUpper not in poolingTypes:
+            raise ValueError(
+                f"Invalid pooling type: {poolingType}. "
+                + f"Valid values are: {poolingTypes}"
+            )
+        return self._set(poolingType=poolingType)
+
+    def setModelDraft(self, modelDraft: str):
+        """Set the draft model for speculative decoding"""
+        return self._set(modelDraft=modelDraft)
+
+    def setModelAlias(self, modelAlias: str):
+        """Set a model alias"""
+        return self._set(modelAlias=modelAlias)
+
+    def setLookupCacheStaticFilePath(self, lookupCacheStaticFilePath: str):
+        """Set path to static lookup cache to use for lookup decoding (not updated by generation)"""
+        return self._set(lookupCacheStaticFilePath=lookupCacheStaticFilePath)
+
+    def setLookupCacheDynamicFilePath(self, lookupCacheDynamicFilePath: str):
+        """Set path to dynamic lookup cache to use for lookup decoding (updated by generation)"""
+        return self._set(lookupCacheDynamicFilePath=lookupCacheDynamicFilePath)
+
+    def setEmbedding(self, embedding: bool):
+        """Whether to load model with embedding support"""
+        return self._set(embedding=embedding)
+
+    def setFlashAttention(self, flashAttention: bool):
+        """Whether to enable Flash Attention"""
+        return self._set(flashAttention=flashAttention)
+
+    def setInputPrefixBos(self, inputPrefixBos: bool):
+        """Whether to add prefix BOS to user inputs, preceding the `--in-prefix` bool"""
+        return self._set(inputPrefixBos=inputPrefixBos)
+
+    def setUseMmap(self, useMmap: bool):
+        """Whether to use memory-map model (faster load but may increase pageouts if not using mlock)"""
+        return self._set(useMmap=useMmap)
+
+    def setUseMlock(self, useMlock: bool):
+        """Whether to force the system to keep model in RAM rather than swapping or compressing"""
+        return self._set(useMlock=useMlock)
+
+    def setNoKvOffload(self, noKvOffload: bool):
+        """Whether to disable KV offload"""
+        return self._set(noKvOffload=noKvOffload)
+
+    def setSystemPrompt(self, systemPrompt: bool):
+        """Set a system prompt to use"""
+        return self._set(systemPrompt=systemPrompt)
+
+    def setChatTemplate(self, chatTemplate: str):
+        """The chat template to use"""
+        return self._set(chatTemplate=chatTemplate)
+
+    # -------- INFERENCE SETTERS --------
+    def setInputPrefix(self, inputPrefix: str):
+        """Set the prompt to start generation with"""
+        return self._set(inputPrefix=inputPrefix)
 
-def setBeamSize(self, value):
-    """Sets the number of beam size for beam search.
+    def setInputSuffix(self, inputSuffix: str):
+        """Set a suffix for infilling"""
+        return self._set(inputSuffix=inputSuffix)
 
-    Parameters
-    ----------
-    value : int
-        Number of beam size for beam search
-    """
-    return self._set(beamSize=value)
+    def setCachePrompt(self, cachePrompt: bool):
+        """Whether to remember the prompt to avoid reprocessing it"""
+        return self._set(cachePrompt=cachePrompt)
 
+    def setNPredict(self, nPredict: int):
+        """Set the number of tokens to predict"""
+        return self._set(nPredict=nPredict)
 
-def setNReturnSequences(self, value):
-    """Sets the number of sequences to return from the beam search.
+    def setTopK(self, topK: int):
+        """Set top-k sampling"""
+        return self._set(topK=topK)
 
-    Parameters
-    ----------
-    value : int
-        Number of sequences to return
-    """
-    return self._set(nReturnSequences=value)
+    def setTopP(self, topP: float):
+        """Set top-p sampling"""
+        return self._set(topP=topP)
+
+    def setMinP(self, minP: float):
+        """Set min-p sampling"""
+        return self._set(minP=minP)
+
+    def setTfsZ(self, tfsZ: float):
+        """Set tail free sampling, parameter z"""
+        return self._set(tfsZ=tfsZ)
+
+    def setTypicalP(self, typicalP: float):
+        """Set locally typical sampling, parameter p"""
+        return self._set(typicalP=typicalP)
+
+    def setTemperature(self, temperature: float):
+        """Set the temperature"""
+        return self._set(temperature=temperature)
+
+    def setDynamicTemperatureRange(self, dynamicTemperatureRange: float):
+        """Set the dynamic temperature range"""
+        return self._set(dynamicTemperatureRange=dynamicTemperatureRange)
+
+    def setDynamicTemperatureExponent(self, dynamicTemperatureExponent: float):
+        """Set the dynamic temperature exponent"""
+        return self._set(dynamicTemperatureExponent=dynamicTemperatureExponent)
+
+    def setRepeatLastN(self, repeatLastN: int):
+        """Set the last n tokens to consider for penalties"""
+        return self._set(repeatLastN=repeatLastN)
+
+    def setRepeatPenalty(self, repeatPenalty: float):
+        """Set the penalty of repeated sequences of tokens"""
+        return self._set(repeatPenalty=repeatPenalty)
+
+    def setFrequencyPenalty(self, frequencyPenalty: float):
+        """Set the repetition alpha frequency penalty"""
+        return self._set(frequencyPenalty=frequencyPenalty)
+
+    def setPresencePenalty(self, presencePenalty: float):
+        """Set the repetition alpha presence penalty"""
+        return self._set(presencePenalty=presencePenalty)
+
+    def setMiroStat(self, miroStat: str):
+        """Set MiroStat sampling strategies."""
+        return self._set(miroStat=miroStat)
+
+    def setMiroStatTau(self, miroStatTau: float):
+        """Set the MiroStat target entropy, parameter tau"""
+        return self._set(miroStatTau=miroStatTau)
+
+    def setMiroStatEta(self, miroStatEta: float):
+        """Set the MiroStat learning rate, parameter eta"""
+        return self._set(miroStatEta=miroStatEta)
+
+    def setPenalizeNl(self, penalizeNl: bool):
+        """Whether to penalize newline tokens"""
+        return self._set(penalizeNl=penalizeNl)
+
+    def setNKeep(self, nKeep: int):
+        """Set the number of tokens to keep from the initial prompt"""
+        return self._set(nKeep=nKeep)
+
+    def setSeed(self, seed: int):
+        """Set the RNG seed"""
+        return self._set(seed=seed)
+
+    def setNProbs(self, nProbs: int):
+        """Set the amount top tokens probabilities to output if greater than 0."""
+        return self._set(nProbs=nProbs)
+
+    def setMinKeep(self, minKeep: int):
+        """Set the amount of tokens the samplers should return at least (0 = disabled)"""
+        return self._set(minKeep=minKeep)
+
+    def setGrammar(self, grammar: bool):
+        """Set BNF-like grammar to constrain generations"""
+        return self._set(grammar=grammar)
+
+    def setPenaltyPrompt(self, penaltyPrompt: str):
+        """Override which part of the prompt is penalized for repetition."""
+        return self._set(penaltyPrompt=penaltyPrompt)
+
+    def setIgnoreEos(self, ignoreEos: bool):
+        """Set whether to ignore end of stream token and continue generating (implies --logit-bias 2-inf)"""
+        return self._set(ignoreEos=ignoreEos)
+
+    def setDisableTokenIds(self, disableTokenIds: List[int]):
+        """Set the token ids to disable in the completion"""
+        return self._set(disableTokenIds=disableTokenIds)
+
+    def setStopStrings(self, stopStrings: List[str]):
+        """Set strings upon seeing which token generation is stopped"""
+        return self._set(stopStrings=stopStrings)
+
+    def setSamplers(self, samplers: List[str]):
+        """Set which samplers to use for token generation in the given order"""
+        return self._set(samplers=samplers)
+
+    def setUseChatTemplate(self, useChatTemplate: bool):
+        """Set whether generate should apply a chat template"""
+        return self._set(useChatTemplate=useChatTemplate)
+    
+    def setNParallel(self, nParallel: int):
+        """Sets the number of parallel processes for decoding. This is an alias for `setBatchSize`."""
+        return self.setBatchSize(nParallel)
+
+    # -------- JAVA SETTERS --------
+    def setTokenIdBias(self, tokenIdBias: Dict[int, float]):
+        """Set token id bias"""
+        return self._call_java("setTokenIdBias", tokenIdBias)
+
+    def setTokenBias(self, tokenBias: Dict[str, float]):
+        """Set token id bias"""
+        return self._call_java("setTokenBias", tokenBias)
+
+    def setLoraAdapters(self, loraAdapters: Dict[str, float]):
+        """Set token id bias"""
+        return self._call_java("setLoraAdapters", loraAdapters)
+
+    def getMetadata(self):
+        """Gets the metadata of the model"""
+        return self._call_java("getMetadata")