edsnlp.processing

helpers

DataFrames = None module-attribute

spec = importlib.util.find_spec(module.value) module-attribute

DataFrameModules

Bases: Enum

Source code in edsnlp/processing/helpers.py

class DataFrameModules(Enum):
    PANDAS = "pandas"
    PYSPARK = "pyspark.sql"
    KOALAS = "databricks.koalas"

PANDAS = 'pandas' class-attribute

PYSPARK = 'pyspark.sql' class-attribute

KOALAS = 'databricks.koalas' class-attribute

get_module(df)

Source code in edsnlp/processing/helpers.py

def get_module(df: DataFrames):
    for module in list(DataFrameModules):
        if df.__class__.__module__.startswith(module.value):
            return module

check_spacy_version_for_context()

Source code in edsnlp/processing/helpers.py

def check_spacy_version_for_context():  # pragma: no cover
    import spacy

    spacy_version = getattr(spacy, "__version__")
    if LooseVersion(spacy_version) < LooseVersion("3.2"):
        raise VersionConflict(
            "You provided a `context` argument, which only work with spacy>=3.2.\n"
            f"However, we found SpaCy version {spacy_version}.\n",
            "Please upgrade SpaCy ;)",
        )

simple

nlp = spacy.blank('fr') module-attribute

ExtensionSchema = Union[str, List[str], Dict[str, Any]] module-attribute

_df_to_spacy(note, nlp, context)

Takes a pandas DataFrame and return a generator that can be used in nlp.pipe().

PARAMETER	DESCRIPTION
note	A pandas DataFrame with at least note_text and note_id columns. A Doc object will be created for each line. TYPE: pd.DataFrame

RETURNS	DESCRIPTION
generator	A generator which items are of the form (text, context), with text being a string and context a dictionnary

Source code in edsnlp/processing/simple.py

def _df_to_spacy(
    note: pd.DataFrame,
    nlp: Language,
    context: List[str],
):
    """
    Takes a pandas DataFrame and return a generator that can be used in
    `nlp.pipe()`.

    Parameters
    ----------
    note: pd.DataFrame
        A pandas DataFrame with at least `note_text` and `note_id` columns.
        A `Doc` object will be created for each line.

    Returns
    -------
    generator:
        A generator which items are of the form (text, context), with `text`
        being a string and `context` a dictionnary
    """

    if context:
        check_spacy_version_for_context()

    kept_cols = ["note_text"] + context

    for col in kept_cols:
        if col not in note.columns:
            raise ValueError(f"No column named {repr(col)} found in df")

    def add_context(context_values):
        note_text = context_values.note_text
        doc = nlp.make_doc(note_text)
        for col in context:
            doc._.set(col, getattr(context_values, col))
        return doc

    yield from map(
        add_context,
        note[kept_cols].itertuples(),
    )

_flatten(list_of_lists)

Flatten a list of lists to a combined list.

Source code in edsnlp/processing/simple.py

def _flatten(list_of_lists: List[List[Any]]):
    """
    Flatten a list of lists to a combined list.
    """
    return [item for sublist in list_of_lists for item in sublist]

_pipe_generator(note, nlp, context=[], additional_spans='discarded', extensions=[], batch_size=50, progress_bar=True)

Source code in edsnlp/processing/simple.py

def _pipe_generator(
    note: pd.DataFrame,
    nlp: Language,
    context: List[str] = [],
    additional_spans: Union[List[str], str] = "discarded",
    extensions: ExtensionSchema = [],
    batch_size: int = 50,
    progress_bar: bool = True,
):

    if type(extensions) == str:
        extensions = [extensions]
    elif type(extensions) == dict:
        extensions = list(extensions.keys())

    if type(additional_spans) == str:
        additional_spans = [additional_spans]

    if "note_id" not in context:
        context.append("note_id")

    if not nlp.has_pipe("eds.context"):
        nlp.add_pipe("eds.context", first=True, config=dict(context=context))

    gen = _df_to_spacy(note, nlp, context)
    n_docs = len(note)
    pipeline = nlp.pipe(gen, batch_size=batch_size)

    for doc in tqdm(pipeline, total=n_docs, disable=not progress_bar):

        yield _full_schema(
            doc,
            additional_spans=additional_spans,
            extensions=extensions,
        )

_single_schema(ent, span_type='ents', extensions=[])

Source code in edsnlp/processing/simple.py

def _single_schema(
    ent: Span,
    span_type: str = "ents",
    extensions: List[str] = [],
):

    return {
        "note_id": ent.doc._.note_id,
        "lexical_variant": ent.text,
        "label": ent.label_,
        "span_type": span_type,
        "start": ent.start_char,
        "end": ent.end_char,
        **{extension: getattr(ent._, extension) for extension in extensions},
    }

_full_schema(doc, additional_spans=[], extensions=[])

Function used when Parallelising tasks via joblib. Takes a Doc as input, and returns a list of serializable objects

Note

The parallelisation needs for output objects to be serializable: after splitting the task into separate jobs, intermediate results are saved on memory before being aggregated, thus the need to be serializable. For instance, spaCy's spans aren't serializable since they are merely a view of the parent document.

Check the source code of this function for an example.

Source code in edsnlp/processing/simple.py

def _full_schema(
    doc: Doc,
    additional_spans: List[str] = [],
    extensions: List[str] = [],
):
    """
    Function used when Parallelising tasks via joblib.
    Takes a Doc as input, and returns a list of serializable objects

    !!! note

        The parallelisation needs for output objects to be **serializable**:
        after splitting the task into separate jobs, intermediate results
        are saved on memory before being aggregated, thus the need to be
        serializable. For instance, spaCy's spans aren't serializable since
        they are merely a *view* of the parent document.

        Check the source code of this function for an example.

    """

    results = []

    results.extend(
        [
            _single_schema(
                ent,
                extensions=extensions,
            )
            for ent in doc.ents
            if doc.ents
        ]
    )

    for span_type in additional_spans:
        results.extend(
            [
                _single_schema(
                    ent,
                    span_type=span_type,
                    extensions=extensions,
                )
                for ent in doc.spans[span_type]
                if doc.spans[span_type]
            ]
        )
    return results

pipe(note, nlp, context=[], additional_spans='discarded', extensions=[], batch_size=1000, progress_bar=True)

Function to apply a spaCy pipe to a pandas DataFrame note For a large DataFrame, prefer the parallel version.

PARAMETER	DESCRIPTION
note	A pandas DataFrame with a note_id and note_text column TYPE: DataFrame
nlp	A spaCy pipe TYPE: Language
context	A list of column to add to the generated SpaCy document as an extension. For instance, if context=["note_datetime"], the corresponding value found in thenote_datetimecolumn will be stored indoc._.note_datetime, which can be useful e.g. for thedates` pipeline. TYPE: List[str] DEFAULT: []
additional_spans	A name (or list of names) of SpanGroup on which to apply the pipe too: SpanGroup are available as doc.spans[spangroup_name] and can be generated by some pipes. For instance, the date pipe populates doc.spans['dates'] TYPE: Union[List[str], str], by default "discarded" DEFAULT: 'discarded'
extensions	Spans extensions to add to the extracted results: For instance, if extensions=["score_name"], the extracted result will include, for each entity, ent._.score_name. TYPE: List[Tuple[str, T.DataType]], by default [] DEFAULT: []
batch_size	Batch size used by spaCy's pipe TYPE: int, by default 1000 DEFAULT: 1000
progress_bar	Whether to display a progress bar or not TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
DataFrame	A pandas DataFrame with one line per extraction

Source code in edsnlp/processing/simple.py

def pipe(
    note: pd.DataFrame,
    nlp: Language,
    context: List[str] = [],
    additional_spans: Union[List[str], str] = "discarded",
    extensions: Union[List[str], str] = [],
    batch_size: int = 1000,
    progress_bar: bool = True,
):
    """
    Function to apply a spaCy pipe to a pandas DataFrame note
    For a large DataFrame, prefer the parallel version.

    Parameters
    ----------
    note : DataFrame
        A pandas DataFrame with a `note_id` and `note_text` column
    nlp : Language
        A spaCy pipe
    context : List[str]
        A list of column to add to the generated SpaCy document as an extension.
        For instance, if `context=["note_datetime"], the corresponding value found
        in the `note_datetime` column will be stored in `doc._.note_datetime`,
        which can be useful e.g. for the `dates` pipeline.
    additional_spans : Union[List[str], str], by default "discarded"
        A name (or list of names) of SpanGroup on which to apply the pipe too:
        SpanGroup are available as `doc.spans[spangroup_name]` and can be generated
        by some pipes. For instance, the `date` pipe populates doc.spans['dates']
    extensions : List[Tuple[str, T.DataType]], by default []
        Spans extensions to add to the extracted results:
        For instance, if `extensions=["score_name"]`, the extracted result
        will include, for each entity, `ent._.score_name`.
    batch_size : int, by default 1000
        Batch size used by spaCy's pipe
    progress_bar: bool, by default True
        Whether to display a progress bar or not

    Returns
    -------
    DataFrame
        A pandas DataFrame with one line per extraction
    """
    return pd.DataFrame(
        _flatten(
            _pipe_generator(
                note=note,
                nlp=nlp,
                context=context,
                additional_spans=additional_spans,
                extensions=extensions,
                batch_size=batch_size,
                progress_bar=progress_bar,
            )
        )
    )

wrapper

pipe(note, nlp, n_jobs=-2, context=[], additional_spans='discarded', extensions=[], **kwargs)

Function to apply a spaCy pipe to a pandas or pyspark DataFrame

PARAMETER	DESCRIPTION
note	A pandas/pyspark/koalas DataFrame with a note_id and note_text column TYPE: DataFrame
nlp	A spaCy pipe TYPE: Language
context	A list of column to add to the generated SpaCy document as an extension. For instance, if context=["note_datetime"], the corresponding value found in thenote_datetimecolumn will be stored indoc._.note_datetime, which can be useful e.g. for thedates` pipeline. TYPE: List[str] DEFAULT: []
n_jobs	Only used when providing a Pandas DataFrame n_jobs=1 corresponds to simple_pipe n_jobs>1 corresponds to parallel_pipe with n_jobs parallel workers n_jobs=-1 corresponds to parallel_pipe with maximun number of workers n_jobs=-2 corresponds to parallel_pipe with maximun number of workers -1 TYPE: int, by default -2 DEFAULT: -2
additional_spans	A name (or list of names) of SpanGroup on which to apply the pipe too: SpanGroup are available as doc.spans[spangroup_name] and can be generated by some pipes. For instance, the date pipe populates doc.spans['dates'] TYPE: Union[List[str], str], by default "discarded" DEFAULT: 'discarded'
extensions	Spans extensions to add to the extracted results: For instance, if extensions=["score_name"], the extracted result will include, for each entity, ent._.score_name. TYPE: List[Tuple[str, T.DataType]], by default [] DEFAULT: []
kwargs	Additional parameters depending on the how argument. TYPE: Dict[str, Any]

RETURNS	DESCRIPTION
DataFrame	A DataFrame with one line per extraction

Source code in edsnlp/processing/wrapper.py

def pipe(
    note: DataFrames,
    nlp: Language,
    n_jobs: int = -2,
    context: List[str] = [],
    additional_spans: Union[List[str], str] = "discarded",
    extensions: ExtensionSchema = [],
    **kwargs: Dict[str, Any],
) -> DataFrames:
    """
    Function to apply a spaCy pipe to a pandas or pyspark DataFrame


    Parameters
    ----------
    note : DataFrame
        A pandas/pyspark/koalas DataFrame with a `note_id` and `note_text` column
    nlp : Language
        A spaCy pipe
    context : List[str]
        A list of column to add to the generated SpaCy document as an extension.
        For instance, if `context=["note_datetime"], the corresponding value found
        in the `note_datetime` column will be stored in `doc._.note_datetime`,
        which can be useful e.g. for the `dates` pipeline.
    n_jobs : int, by default -2
        Only used when providing a Pandas DataFrame

        - `n_jobs=1` corresponds to `simple_pipe`
        - `n_jobs>1` corresponds to `parallel_pipe` with `n_jobs` parallel workers
        - `n_jobs=-1` corresponds to `parallel_pipe` with maximun number of workers
        - `n_jobs=-2` corresponds to `parallel_pipe` with maximun number of workers -1
    additional_spans : Union[List[str], str], by default "discarded"
        A name (or list of names) of SpanGroup on which to apply the pipe too:
        SpanGroup are available as `doc.spans[spangroup_name]` and can be generated
        by some pipes. For instance, the `date` pipe populates doc.spans['dates']
    extensions : List[Tuple[str, T.DataType]], by default []
        Spans extensions to add to the extracted results:
        For instance, if `extensions=["score_name"]`, the extracted result
        will include, for each entity, `ent._.score_name`.
    kwargs : Dict[str, Any]
        Additional parameters depending on the `how` argument.

    Returns
    -------
    DataFrame
        A DataFrame with one line per extraction
    """

    module = get_module(note)

    if module == DataFrameModules.PANDAS:
        if n_jobs == 1:

            return simple_pipe(
                note=note,
                nlp=nlp,
                context=context,
                additional_spans=additional_spans,
                extensions=extensions,
                **kwargs,
            )

        else:

            return parallel_pipe(
                note=note,
                nlp=nlp,
                context=context,
                additional_spans=additional_spans,
                extensions=extensions,
                n_jobs=n_jobs,
                **kwargs,
            )

    if extensions and type(extensions) != dict:
        raise ValueError(
            """
            When using Spark or Koalas, you should provide extension names
            along with the extension type (as a dictionnary):
            `d[extension_name] = extension_type`
            """  # noqa W291
        )

    from .distributed import pipe as distributed_pipe

    return distributed_pipe(
        note=note,
        nlp=nlp,
        context=context,
        additional_spans=additional_spans,
        extensions=extensions,
        **kwargs,
    )

parallel

nlp = spacy.blank('fr') module-attribute

_define_nlp(new_nlp)

Set the global nlp variable Doing it this way saves non negligeable amount of time

Source code in edsnlp/processing/parallel.py

def _define_nlp(new_nlp: Language):
    """
    Set the global nlp variable
    Doing it this way saves non negligeable amount of time
    """
    global nlp
    nlp = new_nlp

_chunker(iterable, total_length, chunksize)

Takes an iterable and chunk it.

Source code in edsnlp/processing/parallel.py

def _chunker(
    iterable: Iterable,
    total_length: int,
    chunksize: int,
):
    """
    Takes an iterable and chunk it.
    """
    return (
        iterable[pos : pos + chunksize] for pos in range(0, total_length, chunksize)
    )

_process_chunk(note, **pipe_kwargs)

Source code in edsnlp/processing/parallel.py

def _process_chunk(note: pd.DataFrame, **pipe_kwargs):

    list_results = []

    for out in _pipe_generator(note, nlp, progress_bar=False, **pipe_kwargs):
        list_results += out

    return list_results

pipe(note, nlp, context=[], additional_spans='discarded', extensions=[], chunksize=100, n_jobs=-2, progress_bar=True, **pipe_kwargs)

Function to apply a spaCy pipe to a pandas DataFrame note by using multiprocessing

PARAMETER	DESCRIPTION
note	A pandas DataFrame with a note_id and note_text column TYPE: DataFrame
nlp	A spaCy pipe TYPE: Language
context	A list of column to add to the generated SpaCy document as an extension. For instance, if context=["note_datetime"], the corresponding value found in thenote_datetimecolumn will be stored indoc._.note_datetime, which can be useful e.g. for thedates` pipeline. TYPE: List[str] DEFAULT: []
additional_spans	A name (or list of names) of SpanGroup on which to apply the pipe too: SpanGroup are available as doc.spans[spangroup_name] and can be generated by some pipes. For instance, the date pipe populates doc.spans['dates'] TYPE: Union[List[str], str], by default "discarded" DEFAULT: 'discarded'
extensions	Spans extensions to add to the extracted results: FOr instance, if extensions=["score_name"], the extracted result will include, for each entity, ent._.score_name. TYPE: List[Tuple[str, T.DataType]], by default [] DEFAULT: []
chunksize	Batch size used to split tasks TYPE: int DEFAULT: 100
n_jobs	Max number of parallel jobs. The default value uses the maximum number of available cores. TYPE: int DEFAULT: -2
progress_bar	Whether to display a progress bar or not TYPE: bool DEFAULT: True
**pipe_kwargs	Arguments exposed in processing.pipe_generator are also available here DEFAULT: {}

RETURNS	DESCRIPTION
DataFrame	A pandas DataFrame with one line per extraction

Source code in edsnlp/processing/parallel.py

def pipe(
    note: pd.DataFrame,
    nlp: Language,
    context: List[str] = [],
    additional_spans: Union[List[str], str] = "discarded",
    extensions: ExtensionSchema = [],
    chunksize: int = 100,
    n_jobs: int = -2,
    progress_bar: bool = True,
    **pipe_kwargs,
):
    """
    Function to apply a spaCy pipe to a pandas DataFrame note by using multiprocessing

    Parameters
    ----------
    note : DataFrame
        A pandas DataFrame with a `note_id` and `note_text` column
    nlp : Language
        A spaCy pipe
    context : List[str]
        A list of column to add to the generated SpaCy document as an extension.
        For instance, if `context=["note_datetime"], the corresponding value found
        in the `note_datetime` column will be stored in `doc._.note_datetime`,
        which can be useful e.g. for the `dates` pipeline.
    additional_spans : Union[List[str], str], by default "discarded"
        A name (or list of names) of SpanGroup on which to apply the pipe too:
        SpanGroup are available as `doc.spans[spangroup_name]` and can be generated
        by some pipes. For instance, the `date` pipe populates doc.spans['dates']
    extensions : List[Tuple[str, T.DataType]], by default []
        Spans extensions to add to the extracted results:
        FOr instance, if `extensions=["score_name"]`, the extracted result
        will include, for each entity, `ent._.score_name`.
    chunksize: int, by default 100
        Batch size used to split tasks
    n_jobs: int, by default -2
        Max number of parallel jobs.
        The default value uses the maximum number of available cores.
    progress_bar: bool, by default True
        Whether to display a progress bar or not
    **pipe_kwargs:
        Arguments exposed in `processing.pipe_generator` are also available here

    Returns
    -------
    DataFrame
        A pandas DataFrame with one line per extraction
    """

    if context:
        check_spacy_version_for_context()

    # Setting the nlp variable
    _define_nlp(nlp)

    verbose = 10 if progress_bar else 0

    executor = Parallel(
        n_jobs, backend="multiprocessing", prefer="processes", verbose=verbose
    )
    executor.warn(f"Used nlp components: {nlp.component_names}")

    pipe_kwargs["additional_spans"] = additional_spans
    pipe_kwargs["extensions"] = extensions
    pipe_kwargs["context"] = context

    if verbose:
        executor.warn(f"{int(len(note)/chunksize)} tasks to complete")

    do = delayed(_process_chunk)

    tasks = (
        do(chunk, **pipe_kwargs)
        for chunk in _chunker(note, len(note), chunksize=chunksize)
    )
    result = executor(tasks)

    out = _flatten(result)

    return pd.DataFrame(out)

distributed

pyspark_type_finder(obj)

Returns (when possible) the PySpark type of any python object

Source code in edsnlp/processing/distributed.py

def pyspark_type_finder(obj):
    """
    Returns (when possible) the PySpark type of any python object
    """
    try:
        inferred_type = T._infer_type(obj)
        print(f"Inferred type is {repr(inferred_type)}")
        return inferred_type
    except TypeError:
        raise TypeError("Cannot infer type for this object.")

module_checker(func, *args, **kwargs)

Source code in edsnlp/processing/distributed.py

@decorator
def module_checker(
    func: Callable,
    *args,
    **kwargs,
) -> Any:

    args = list(args)
    note = args.pop(0)
    module = get_module(note)

    if module == DataFrameModules.PYSPARK:
        return func(note, *args, **kwargs)
    elif module == DataFrameModules.KOALAS:
        import databricks.koalas  # noqa F401

        note_spark = note.to_spark()
        note_nlp_spark = func(note_spark, *args, **kwargs)
        return note_nlp_spark.to_koalas()

pipe(note, nlp, context=[], additional_spans='discarded', extensions=[])

Function to apply a spaCy pipe to a pyspark or koalas DataFrame note

PARAMETER	DESCRIPTION
note	A Pyspark or Koalas DataFrame with a note_id and note_text column TYPE: DataFrame
nlp	A spaCy pipe TYPE: Language
context	A list of column to add to the generated SpaCy document as an extension. For instance, if context=["note_datetime"], the corresponding value found in thenote_datetimecolumn will be stored indoc._.note_datetime, which can be useful e.g. for thedates` pipeline. TYPE: List[str] DEFAULT: []
additional_spans	A name (or list of names) of SpanGroup on which to apply the pipe too: SpanGroup are available as doc.spans[spangroup_name] and can be generated by some pipes. For instance, the date pipe populates doc.spans['dates'] TYPE: Union[List[str], str], by default "discarded" DEFAULT: 'discarded'
extensions	Spans extensions to add to the extracted results: FOr instance, if extensions=["score_name"], the extracted result will include, for each entity, ent._.score_name. TYPE: List[Tuple[str, T.DataType]], by default [] DEFAULT: []

RETURNS	DESCRIPTION
DataFrame	A pyspark DataFrame with one line per extraction

Source code in edsnlp/processing/distributed.py

@module_checker
def pipe(
    note: DataFrames,
    nlp: Language,
    context: List[str] = [],
    additional_spans: Union[List[str], str] = "discarded",
    extensions: List[Tuple[str, T.DataType]] = [],
) -> DataFrame:
    """
    Function to apply a spaCy pipe to a pyspark or koalas DataFrame note

    Parameters
    ----------
    note : DataFrame
        A Pyspark or Koalas DataFrame with a `note_id` and `note_text` column
    nlp : Language
        A spaCy pipe
    context : List[str]
        A list of column to add to the generated SpaCy document as an extension.
        For instance, if `context=["note_datetime"], the corresponding value found
        in the `note_datetime` column will be stored in `doc._.note_datetime`,
        which can be useful e.g. for the `dates` pipeline.
    additional_spans : Union[List[str], str], by default "discarded"
        A name (or list of names) of SpanGroup on which to apply the pipe too:
        SpanGroup are available as `doc.spans[spangroup_name]` and can be generated
        by some pipes. For instance, the `date` pipe populates doc.spans['dates']
    extensions : List[Tuple[str, T.DataType]], by default []
        Spans extensions to add to the extracted results:
        FOr instance, if `extensions=["score_name"]`, the extracted result
        will include, for each entity, `ent._.score_name`.

    Returns
    -------
    DataFrame
        A pyspark DataFrame with one line per extraction
    """

    if context:
        check_spacy_version_for_context()

    spark = SparkSession.builder.enableHiveSupport().getOrCreate()
    sc = spark.sparkContext

    if not nlp.has_pipe("eds.context"):
        nlp.add_pipe("eds.context", first=True, config=dict(context=context))

    nlp_bc = sc.broadcast(nlp)

    def _udf_factory(
        additional_spans: Union[List[str], str] = "discarded",
        extensions: Dict[str, T.DataType] = dict(),
    ):

        schema = T.ArrayType(
            T.StructType(
                [
                    T.StructField("lexical_variant", T.StringType(), False),
                    T.StructField("label", T.StringType(), False),
                    T.StructField("span_type", T.StringType(), True),
                    T.StructField("start", T.IntegerType(), False),
                    T.StructField("end", T.IntegerType(), False),
                    *[
                        T.StructField(extension_name, extension_type, True)
                        for extension_name, extension_type in extensions.items()
                    ],
                ]
            )
        )

        def f(
            text,
            *context_values,
            additional_spans=additional_spans,
            extensions=extensions,
        ):

            if text is None:
                return []

            nlp = nlp_bc.value

            for _, pipe in nlp.pipeline:
                if isinstance(pipe, BaseComponent):
                    pipe.set_extensions()

            doc = nlp.make_doc(text)
            for context_name, context_value in zip(context, context_values):
                doc._.set(context_name, context_value)
            doc = nlp(doc)

            ents = []

            for ent in doc.ents:
                parsed_extensions = [
                    getattr(ent._, extension) for extension in extensions.keys()
                ]

                ents.append(
                    (
                        ent.text,
                        ent.label_,
                        "ents",
                        ent.start_char,
                        ent.end_char,
                        *parsed_extensions,
                    )
                )

            if additional_spans is None:
                return ents

            if type(additional_spans) == str:
                additional_spans = [additional_spans]

            for spans_name in additional_spans:

                for ent in doc.spans.get(spans_name, []):

                    parsed_extensions = [
                        getattr(ent._, extension) for extension in extensions.keys()
                    ]

                    ents.append(
                        (
                            ent.text,
                            ent.label_,
                            spans_name,
                            ent.start_char,
                            ent.end_char,
                            *parsed_extensions,
                        )
                    )

            return ents

        f_udf = F.udf(
            partial(
                f,
                additional_spans=additional_spans,
                extensions=extensions,
            ),
            schema,
        )

        return f_udf

    matcher = _udf_factory(
        additional_spans=additional_spans,
        extensions=extensions,
    )

    note_nlp = note.withColumn(
        "matches", matcher(F.col("note_text"), *[F.col(c) for c in context])
    )
    note_nlp = note_nlp.withColumn("matches", F.explode(note_nlp.matches))

    note_nlp = note_nlp.select("note_id", "matches.*")

    return note_nlp