?? KeyBERT-BG - 保加利亚文关键词提取

KeyBERT-BG 是一个用于提取保加利亚文关键词的模型。使用的数据集是自定义的 this ，我已经上传到 Kaggle 上。

使用方法

导入库：

import re
from typing import Dict
from pprint import pprint

from transformers import AutoTokenizer, AutoModelForTokenClassification

首先，您需要定义此方法，因为文本预处理是自定义的，标准管道方法不适用：

def extract_keywords(
    text: str, 
    model_id="auhide/keybert-bg", 
    max_len: int = 300,
    id2group: Dict[int, str] = {
        # Indicates that this is not a keyword.
        0: "O",
        # Begining of keyword.
        1: "B-KWD",
        # Additional keywords (might also indicate the end of a keyword sequence).
        # You can merge these with the begining keyword `B-KWD`.
        2: "I-KWD",
    },
    # Probability threshold based on which the keywords will be accepted.
    # If their probabiliy is less than `threshold`, they won't be added to the list of keywords.
    threshold=0.50
):
    # Initialize the tokenizer and model.
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    keybert = AutoModelForTokenClassification.from_pretrained(model_id)

    # Preprocess the text.
    # Surround punctuation with whitespace and convert multiple whitespaces
    # into single ones.
    text = re.sub(r"([,\.?!;:\'\"\(\)\[\]„”])", r" \1 ", text)
    text = re.sub(r"\s+", r" ", text)
    words = text.split()

    # Tokenize the processed `text` (this includes padding or truncation).
    tokens_data = tokenizer(
        text.strip(), 
        padding="max_length", 
        max_length=max_len, 
        truncation=True, 
        return_tensors="pt"
    )
    input_ids = tokens_data.input_ids
    attention_mask = tokens_data.attention_mask

    # Predict the keywords.
    out = keybert(input_ids, attention_mask=attention_mask).logits
    # Softmax the last dimension so that the probabilities add up to 1.0.
    out = out.softmax(-1)
    # Based on the probabilities, generate the most probable keywords.
    out_argmax = out.argmax(-1)
    prediction = out_argmax.squeeze(0).tolist()
    probabilities = out.squeeze(0)
    
    return [
        {
            # Since the list of words does not have a [CLS] token, the index `i`
            # is one step forward, which means that if we want to access the 
            # appropriate keyword we should use the index `i - 1`.
            "entity": words[i - 1],
            "entity_group": id2group[idx],
            "score": float(probabilities[i, idx])
        } 
        for i, idx in enumerate(prediction) 
        if (idx == 1 or idx == 2) and float(probabilities[i, idx]) > threshold
    ]

选择一段文本并在其上使用该模型。例如，我选择了 this 文章。然后，您可以调用 extract_keywords 提取关键词：

# Reading the text from a file, since it is an article, and the text is large.
with open("input_text.txt", "r", encoding="utf-8") as f:
    text = f.read()

# You can change the threshold based on your needs.
keywords = extract_keywords(text, threshold=0.5)
print("Keywords:")
pprint(keywords)

Keywords:
[{'entity': 'Туитър', 'entity_group': 'B-KWD', 'score': 0.9278278946876526},
 {'entity': 'Илон', 'entity_group': 'B-KWD', 'score': 0.5862686634063721},
 {'entity': 'Мъск', 'entity_group': 'B-KWD', 'score': 0.5289096832275391},
 {'entity': 'изпълнителен',
  'entity_group': 'B-KWD',
  'score': 0.679943323135376},
 {'entity': 'директор', 'entity_group': 'I-KWD', 'score': 0.6161141991615295}]

请注意，您可以使用 transformers 中的 pipeline 方法，但结果会稍差。