examples/tutorial_pyltp/ner.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
https://huggingface.co/LTP
"""
import argparse
import os

from pyltp import Postagger, Segmentor, NamedEntityRecognizer

from project_settings import project_path


def get_args():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--text",
        default="元芳你怎么看？我就趴窗口上看呗！",
        type=str
    )
    parser.add_argument(
        "--ltp_data_dir",
        default=(project_path / "data/pyltp_models/ltp_data_v3.4.0").as_posix(),
        type=str
    )
    args = parser.parse_args()
    return args


def main():
    args = get_args()

    cws_model_path = os.path.join(args.ltp_data_dir, "cws.model")
    pos_model_path = os.path.join(args.ltp_data_dir, "pos.model")
    ner_model_path = os.path.join(args.ltp_data_dir, "ner.model")

    segmentor = Segmentor(cws_model_path)
    postagger = Postagger(pos_model_path)
    recognizer = NamedEntityRecognizer(ner_model_path)

    words = segmentor.segment(args.text)
    postags = postagger.postag(words)
    ner_tags = recognizer.recognize(words, postags)

    print(words)
    print(postags)
    print(ner_tags)

    segmentor.release()
    postagger.release()
    recognizer.release()

    return


if __name__ == "__main__":
    main()