`sparknlp_jsl.annotator.medical_llm.medical_llm_entity_extractor`#

Contains classes for the MedicalLLMEntityExtractor annotator.

Module Contents#

Classes#

MedicalLLMEntityExtractor

End-to-end LLM-based entity extraction for clinical and healthcare text.

class MedicalLLMEntityExtractor(classname='com.johnsnowlabs.nlp.annotators.ner.dl.MedicalLLMEntityExtractor', java_model=None)#

Bases: sparknlp.annotator.LLMEntityExtractor

End-to-end LLM-based entity extraction for clinical and healthcare text.

It performs entity extraction from medical text using Large Language Models (LLMs) in GGUF format with structured JSON output enforced via BNF grammars. The annotator uses string matching to compute accurate character indices for extracted entities.

This annotator follows the LangExtract pattern from Google Research, combining few-shot prompting with constrained generation through llama.cpp BNF grammars to ensure valid JSON output.

The LLM generates responses in this format (enforced by grammar):

{
  "extractions": [
    {"entity": "PROBLEM", "text": "hypertension"},
    {"entity": "TREATMENT", "text": "lisinopril 10mg"}
  ]
}

The model is loaded via MedicalLLMEntityExtractor.pretrained() to download a pretrained JSL clinical model, or MedicalLLMEntityExtractor.loadSavedModel() to load a local GGUF model:

>>> entity_extractor = MedicalLLMEntityExtractor.pretrained() \
...     .setInputCols(["document"]) \
...     .setOutputCol("entities") \
...     .setEntityTypes(["PROBLEM", "TEST", "TREATMENT"])

Input Annotation types	Output Annotation type
`DOCUMENT`	`CHUNK`

Parameters:

promptTemplate (str, optional) – Custom prompt template for entity extraction. Use {entityTypes} and {text} as placeholders.
entityTypes (List[str], optional) – List of entity types to extract (used in prompt), by default [“PROBLEM”, “TEST”, “TREATMENT”]
caseSensitive (bool, optional) – Whether entity matching is case-sensitive, by default False
fewShotExamples (List[Tuple[str, str]], optional) – Few-shot examples as (input, output_json) tuples to guide the model

Examples

>>> import sparknlp
>>> from sparknlp.base import *
>>> from sparknlp_jsl.annotator import *
>>> from pyspark.ml import Pipeline
>>> documentAssembler = DocumentAssembler() \
...     .setInputCol("text") \
...     .setOutputCol("document")
>>> entity_extractor = MedicalLLMEntityExtractor.pretrained() \
...     .setInputCols(["document"]) \
...     .setOutputCol("entities") \
...     .setEntityTypes(["MEDICATION", "DOSAGE", "ROUTE", "FREQUENCY"]) \
...     .setNPredict(500) \
...     .setTemperature(0.1)
>>> pipeline = Pipeline().setStages([documentAssembler, entity_extractor])
>>> data = spark.createDataFrame([["Patient prescribed 500mg amoxicillin PO TID"]]).toDF("text")
>>> result = pipeline.fit(data).transform(data)
>>> result.select("entities.result", "entities.metadata").show(truncate=False)
+------------------------------+--------------------------------+
|result                        |metadata                        |
+------------------------------+--------------------------------+
|[500mg, amoxicillin, PO, TID] |[{entity -> DOSAGE}, ...]       |
+------------------------------+--------------------------------+

See also

LLMEntityExtractor: open-source base class
MedicalLLM: for general-purpose medical LLM inference
MedicalNerModel: for traditional BiLSTM-CRF medical NER
NerConverterInternal: to further process NER results

batchSize#

cachePrompt#

caseSensitive#

chatTemplate#

defragmentationThreshold#

disableLog#

disableTokenIds#

dynamicTemperatureExponent#

dynamicTemperatureRange#

embedding#

entityTypes#

fewShotExamples#

flashAttention#

frequencyPenalty#

getter_attrs = []#

gpuSplitMode#

grammar#

ignoreEos#

inputAnnotatorTypes#

inputCols#

inputPrefix#

inputSuffix#

lazyAnnotator#

logVerbosity#

mainGpu#

minKeep#

minP#

miroStat#

miroStatEta#

miroStatTau#

modelAlias#

modelDraft#

nBatch#

nCtx#

nDraft#

nGpuLayers#

nGpuLayersDraft#

nKeep#

nPredict#

nProbs#

nThreads#

nThreadsBatch#

nUbatch#

name = 'MedicalLLMEntityExtractor'#

noKvOffload#

numaStrategy#

optionalInputAnnotatorTypes = []#

outputAnnotatorType = 'chunk'#

outputCol#

penalizeNl#

penaltyPrompt#

poolingType#

presencePenalty#

promptTemplate#

reasoningBudget#

repeatLastN#

repeatPenalty#

ropeFreqBase#

ropeFreqScale#

ropeScalingType#

samplers#

seed#

stopStrings#

systemPrompt#

temperature#

tfsZ#

topK#

topP#

typicalP#

uid = ''#

useChatTemplate#

useMlock#

useMmap#

yarnAttnFactor#

yarnBetaFast#

yarnBetaSlow#

yarnExtFactor#

yarnOrigCtx#

clear(param: pyspark.ml.param.Param) → None#: Clears a param from the param map if it has been explicitly set.

close()#: Closes the underlying llama.cpp model backend freeing resources.

copy(extra: pyspark.ml._typing.ParamMap | None = None) → JP#

Creates a copy of this instance with the same uid and some extra params. This implementation first calls Params.copy and then make a copy of the companion Java pipeline component with extra params. So both the Python wrapper and the Java pipeline component get copied.

Parameters:: extra (dict, optional) – Extra parameters to copy to the new instance
Returns:: Copy of this instance
Return type:: JavaParams

explainParam(param: str | Param) → str#: Explains a single param and returns its name, doc, and optional default value and user-supplied value in a string.

explainParams() → str#: Returns the documentation of all params with their optionally default values and user-supplied values.

extractParamMap(extra: pyspark.ml._typing.ParamMap | None = None) → pyspark.ml._typing.ParamMap#

Extracts the embedded default param values and user-supplied values, and then merges them with extra values from input into a flat param map, where the latter value is used if there exist conflicts, i.e., with ordering: default param values < user-supplied values < extra.

Parameters:: extra (dict, optional) – extra param values
Returns:: merged param map
Return type:: dict

getBatchSize()#

Gets current batch size.

Returns:: Current batch size
Return type:: int

getCaseSensitive()#: Get whether entity matching is case-sensitive.

getEntityTypes()#: Get the list of entity types to extract.

getFewShotExamples()#: Get the few-shot examples.

getInputCols()#: Gets current column names of input annotations.

getLazyAnnotator()#: Gets whether Annotator should be evaluated lazily in a RecursivePipeline.

getMetadata()#: Gets the metadata of the model

getOrDefault(param: str) → Any#
getOrDefault(param: Param[T]) → T: Gets the value of a param in the user-supplied param map or its default value. Raises an error if neither is set.

getOutputCol()#: Gets output column name of annotations.

getParam(paramName: str) → Param#: Gets a param by its name.

getParamValue(paramName)#

Gets the value of a parameter.

Parameters:: paramName (str) – Name of the parameter

getPromptTemplate()#: Get the custom prompt template for entity extraction.

hasDefault(param: str | Param[Any]) → bool#: Checks whether a param has a default value.

hasParam(paramName: str) → bool#: Tests whether this instance contains a param with a given (string) name.

inputColsValidation(value)#

isDefined(param: str | Param[Any]) → bool#: Checks whether a param is explicitly set by user or has a default value.

isSet(param: str | Param[Any]) → bool#: Checks whether a param is explicitly set by user.

classmethod load(path: str) → RL#: Reads an ML instance from the input path, a shortcut of read().load(path).

classmethod loadSavedModel(path, spark_session)#

Loads a locally saved GGUF model for LLM-based entity extraction.

Parameters:

path (str) – Path to the saved model folder
spark_session (pyspark.sql.SparkSession) – The current SparkSession

Returns:

The restored model

Return type:

MedicalLLMEntityExtractor

classmethod pretrained(name='jsl_medm_q8_v1', lang='en', remote_loc='clinical/models')#

Downloads and loads a pretrained JSL clinical model.

Parameters:

name (str, optional) – Name of the pretrained model, by default “jsl_medm_q8_v1”
lang (str, optional) – Language of the pretrained model, by default “en”
remote_loc (str, optional) – Remote location of the resource, by default “clinical/models”

Returns:

The restored model

Return type:

MedicalLLMEntityExtractor

classmethod read()#: Returns an MLReader instance for this class.

save(path: str) → None#: Save this ML instance to the given path, a shortcut of ‘write().save(path)’.

set(param: Param, value: Any) → None#: Sets a parameter in the embedded param map.

setBatchSize(v)#

Sets batch size.

Parameters:: v (int) – Batch size

setCachePrompt(cachePrompt: bool)#: Whether to remember the prompt to avoid reprocessing it

setCaseSensitive(value)#

Set whether entity matching is case-sensitive.

Parameters:: value (bool) – True for case-sensitive matching, False for case-insensitive
Returns:: The updated model
Return type:: LLMEntityExtractor

setChatTemplate(chatTemplate: str)#: The chat template to use

setDefragmentationThreshold(defragmentationThreshold: float)#: Set the KV cache defragmentation threshold

setDisableLog(disableLog: bool)#: Whether to disable logging

setDisableTokenIds(disableTokenIds: List[int])#: Set the token ids to disable in the completion

setDynamicTemperatureExponent(dynamicTemperatureExponent: float)#: Set the dynamic temperature exponent

setDynamicTemperatureRange(dynamicTemperatureRange: float)#: Set the dynamic temperature range

setEntityTypes(value)#

Set the list of entity types to extract.

Parameters:: value (List[str]) – List of entity type names
Returns:: The updated model
Return type:: LLMEntityExtractor

setFewShotExamples(value)#

Set few-shot examples to guide the model.

Parameters:: value (List[Tuple[str, str]]) – List of (input_text, json_output) tuples as examples
Returns:: The updated model
Return type:: LLMEntityExtractor

setFlashAttention(flashAttention: bool)#: Whether to enable Flash Attention

setFrequencyPenalty(frequencyPenalty: float)#: Set the repetition alpha frequency penalty

setGpuSplitMode(gpuSplitMode: str)#: Set how to split the model across GPUs

setGrammar(grammar: str)#: Set BNF-like grammar to constrain generations

setIgnoreEos(ignoreEos: bool)#: Set whether to ignore end of stream token and continue generating (implies –logit-bias 2-inf)

setInputCols(*value)#

Sets column names of input annotations.

Parameters:: *value (List[str]) – Input columns for the annotator

setInputPrefix(inputPrefix: str)#: Set the prompt to start generation with

setInputSuffix(inputSuffix: str)#: Set a suffix for infilling

setLazyAnnotator(value)#

Sets whether Annotator should be evaluated lazily in a RecursivePipeline.

Parameters:: value (bool) – Whether Annotator should be evaluated lazily in a RecursivePipeline

setLogVerbosity(logVerbosity: int)#: Set the log verbosity level

setMainGpu(mainGpu: int)#: Set the main GPU that is used for scratch and small tensors.

setMinKeep(minKeep: int)#: Set the amount of tokens the samplers should return at least (0 = disabled)

setMinP(minP: float)#: Set min-p sampling

setMiroStat(miroStat: str)#: Set MiroStat sampling strategies.

setMiroStatEta(miroStatEta: float)#: Set the MiroStat learning rate, parameter eta

setMiroStatTau(miroStatTau: float)#: Set the MiroStat target entropy, parameter tau

setModelAlias(modelAlias: str)#: Set a model alias

setModelDraft(modelDraft: str)#: Set the draft model for speculative decoding

setNBatch(nBatch: int)#: Set the logical batch size for prompt processing (must be >=32 to use BLAS)

setNCtx(nCtx: int)#: Set the size of the prompt context

setNDraft(nDraft: int)#: Set the number of tokens to draft for speculative decoding

setNGpuLayers(nGpuLayers: int)#: Set the number of layers to store in VRAM (-1 - use default)

setNGpuLayersDraft(nGpuLayersDraft: int)#: Set the number of layers to store in VRAM for the draft model (-1 - use default)

setNKeep(nKeep: int)#: Set the number of tokens to keep from the initial prompt

setNParallel(nParallel: int)#: Sets the number of parallel processes for decoding. This is an alias for setBatchSize.

setNPredict(nPredict: int)#: Set the number of tokens to predict

setNProbs(nProbs: int)#: Set the amount top tokens probabilities to output if greater than 0.

setNThreads(nThreads: int)#: Set the number of threads to use during generation

setNThreadsBatch(nThreadsBatch: int)#: Set the number of threads to use during batch and prompt processing

setNUbatch(nUbatch: int)#: Set the physical batch size for prompt processing (must be >=32 to use BLAS)

setNoKvOffload(noKvOffload: bool)#: Whether to disable KV offload

setNumaStrategy(numaStrategy: str)#

Set optimization strategies that help on some NUMA systems (if available)

Possible values:

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

setOutputCol(value)#

Sets output column name of annotations.

Parameters:: value (str) – Name of output column

setParamValue(paramName)#

Sets the value of a parameter.

Parameters:: paramName (str) – Name of the parameter

setParams()#

setPenalizeNl(penalizeNl: bool)#: Whether to penalize newline tokens

setPenaltyPrompt(penaltyPrompt: str)#: Override which part of the prompt is penalized for repetition.

setPoolingType(poolingType: str)#

Set the pooling type for embeddings, use model default if unspecified

Possible values:

MEAN: Mean Pooling
CLS: CLS Pooling
LAST: Last token pooling
RANK: For reranked models

setPresencePenalty(presencePenalty: float)#: Set the repetition alpha presence penalty

setPromptTemplate(value)#

Set custom prompt template for entity extraction.

Parameters:: value (str) – Custom prompt template. Use {entityTypes} and {text} as placeholders.
Returns:: The updated model
Return type:: LLMEntityExtractor

setReasoningBudget(reasoningBudget: int)#: Controls the amount of thinking allowed; currently only one of: -1 for unrestricted thinking budget, or 0 to disable thinking (default: -1)

setRepeatLastN(repeatLastN: int)#: Set the last n tokens to consider for penalties

setRepeatPenalty(repeatPenalty: float)#: Set the penalty of repeated sequences of tokens

setRopeFreqBase(ropeFreqBase: float)#: Set the RoPE base frequency, used by NTK-aware scaling

setRopeFreqScale(ropeFreqScale: float)#: Set the RoPE frequency scaling factor, expands context by a factor of 1/N

setRopeScalingType(ropeScalingType: str)#

Set the RoPE frequency scaling method, defaults to linear unless specified by the model.

Possible values:

NONE: Don’t use any scaling
LINEAR: Linear scaling
YARN: YaRN RoPE scaling

setSamplers(samplers: List[str])#: Set which samplers to use for token generation in the given order

setSeed(seed: int)#: Set the RNG seed

setStopStrings(stopStrings: List[str])#: Set strings upon seeing which token generation is stopped

setSystemPrompt(systemPrompt: str)#: Set a system prompt to use

setTemperature(temperature: float)#: Set the temperature

setTfsZ(tfsZ: float)#: Set tail free sampling, parameter z

setTokenBias(tokenBias: Dict[str, float])#: Set token id bias

setTokenIdBias(tokenIdBias: Dict[int, float])#: Set token id bias

setTopK(topK: int)#: Set top-k sampling

setTopP(topP: float)#: Set top-p sampling

setTypicalP(typicalP: float)#: Set locally typical sampling, parameter p

setUseChatTemplate(useChatTemplate: bool)#: Set whether generate should apply a chat template

setUseMlock(useMlock: bool)#: Whether to force the system to keep model in RAM rather than swapping or compressing

setUseMmap(useMmap: bool)#: Whether to use memory-map model (faster load but may increase pageouts if not using mlock)

setYarnAttnFactor(yarnAttnFactor: float)#: Set the YaRN scale sqrt(t) or attention magnitude

setYarnBetaFast(yarnBetaFast: float)#: Set the YaRN low correction dim or beta

setYarnBetaSlow(yarnBetaSlow: float)#: Set the YaRN high correction dim or alpha

setYarnExtFactor(yarnExtFactor: float)#: Set the YaRN extrapolation mix factor

setYarnOrigCtx(yarnOrigCtx: int)#: Set the YaRN original context size of model

transform(dataset: pyspark.sql.dataframe.DataFrame, params: pyspark.ml._typing.ParamMap | None = None) → pyspark.sql.dataframe.DataFrame#

Transforms the input dataset with optional parameters.

New in version 1.3.0.

Parameters:

dataset (pyspark.sql.DataFrame) – input dataset
params (dict, optional) – an optional param map that overrides embedded params.

Returns:

transformed dataset

Return type:

pyspark.sql.DataFrame

write() → JavaMLWriter#: Returns an MLWriter instance for this ML instance.

sparknlp_jsl.annotator.medical_llm.medical_llm_entity_extractor#

Module Contents#

Classes#

`sparknlp_jsl.annotator.medical_llm.medical_llm_entity_extractor`#