# -*- coding: utf-8 -*-
# SPDX-FileCopyrightText: 2016-2024 PyThaiNLP Project
# SPDX-License-Identifier: Apache-2.0
from typing import List
model = None
[docs]def coreference_resolution(
texts: List[str], model_name: str = "han-coref-v1.0", device: str = "cpu"
):
"""
Coreference Resolution
:param List[str] texts: list of texts to apply coreference resolution to
:param str model_name: coreference resolution model
:param str device: device for running coreference resolution model on\
("cpu", "cuda", and others)
:return: List of texts with coreference resolution
:rtype: List[dict]
:Options for model_name:
* *han-coref-v1.0* - (default) Han-Coref: Thai coreference resolution\
by PyThaiNLP v1.0
:Example:
::
from pythainlp.coref import coreference_resolution
print(
coreference_resolution(
["Bill Gates ได้รับวัคซีน COVID-19 เข็มแรกแล้ว ระบุ ผมรู้สึกสบายมาก"]
)
)
# output:
# [
# {'text': 'Bill Gates ได้รับวัคซีน COVID-19 เข็มแรกแล้ว ระบุ ผมรู้สึกสบายมาก',
# 'clusters_string': [['Bill Gates', 'ผม']],
# 'clusters': [[(0, 10), (50, 52)]]}
# ]
"""
global model
if isinstance(texts, str):
texts = [texts]
if model is None and model_name == "han-coref-v1.0":
from pythainlp.coref.han_coref import HanCoref
model = HanCoref(device=device)
if model:
return model.predict(texts)
return [
{"text": text, "clusters_string": [], "clusters": []} for text in texts
]