使用 Polars 在特定数据框列上应用名称实体识别

Question

我想使用类似于以下问题的极坐标将特定函数应用于特定列：

• 在特定数据框列上应用名称实体识别

上面的问题适用于 pandas，我需要很长时间才能在我的计算机上运行它。所以，我想使用极坐标。 从上面的问题来看：

df = pd.DataFrame({'source': ['Paul', 'Paul'],
                   'target': ['GOOGLE', 'Ferrari'],
                   'edge': ['works at', 'drive']
                   })

    source  target  edge
0   Paul    GOOGLE  works at
1   Paul    Ferrari drive

极地的预期结果：

    source  target  edge      Entitiy
0   Paul    GOOGLE  works at  Person
1   Paul    Ferrari drive     Person

!python -m spacy download en_core_web_sm

import spacy
nlp = spacy.load('en_core_web_sm')
df['Entities'] = df['Text'].apply(lambda sent: [(ent.label_) for ent in nlp(sent).ents])  
df['Entities'][1]

如何将带有标签（人）的列添加到带有极坐标的当前数据帧？ 谢谢你。

Answer 1

稍微基于@Luca的答案，您可以将缓存添加到上一级以避免额外的列表理解并直接跳转到实体标签列表：

显示了 Polars 语法，但同样适用于 pandas：

from functools import lru_cache

@lru_cache(2048)  # << size appropriately 
def entity_labels(s: str) -> list:
    return [(ent.label_) for ent in nlp(s).ents]

df.with_columns( 
    pl.col("target").map_elements( 
        function = entity_labels,
        return_dtype = pl.List(pl.String),
    ).alias("labels")
)

Answer 2

您可以使用以下代码在 Polars 中运行应用：

df_pl.with_columns(
    entities = pl.col('target').map_elements(
        lambda sent: [(ent.label_) for ent in nlp(sent).ents])
)

正如 @jqurious 提到的，不应期望这比 Pandas 更快。我进行了几次测试，它花费的时间与 Pandas 相同。

除了@jquurious 的评论之外，如果某些值重复，您可以减少调用 apply 函数的次数。

您可以通过使用 lru_cache 重新定义函数来做到这一点：

from functools import lru_cache
import spacy
import polars as pl

nlp = spacy.load('en_core_web_sm')

@lru_cache(1024)
def cached_nlp(text):
    return nlp(text)

df_pl.with_columns(
    entities = pl.col('target').map_elements(
        lambda sent: [(ent.label_) for ent in cached_nlp(sent).ents])
)

Answer 3

在这种情况下，极地动物将遇到与熊猫相同的问题。

使用

.apply

意味着你本质上使用了 python for 循环。

您可以尝试与多处理池并行运行UDF（用户定义函数）。

根据特定的函数/数据集 - 它可能会也可能不会提供加速，因为多处理本身有其自身的成本 - 必须根据具体情况进行衡量。

在这种情况下，如果我将示例扩展到 10_000 行 - 它的运行速度会快 4 倍。

import spacy
import polars as pl
from   functools import partial
from   multiprocessing import cpu_count, get_context

def recognize(source, nlp):
    return [ent.label_ for ent in nlp(source).ents]

if __name__ == "__main__":
    nlp = spacy.load("en_core_web_sm")

    df = pl.DataFrame({
        "source": ["Paul", "Paul"],
        "target": ["GOOGLE", "Ferrari"],
        "edge": ["works at", "drive"]
    })

    func = partial(recognize, nlp=nlp) # used to pass in `nlp` 
    
    n_workers = cpu_count() // 2 # experiment with value

    with get_context("spawn").Pool(n_workers) as pool:
        df = df.with_columns(Entity = pl.Series(pool.map(func, df.get_column("source"))))
        print(df)

shape: (2, 4)
┌────────┬─────────┬──────────┬────────────┐
│ source | target  | edge     | Entity     │
│ ---    | ---     | ---      | ---        │
│ str    | str     | str      | list[str]  │
╞════════╪═════════╪══════════╪════════════╡
│ Paul   | GOOGLE  | works at | ["PERSON"] │
│ Paul   | Ferrari | drive    | ["PERSON"] │
└────────┴─────────┴──────────┴────────────┘