Rubert Tiny2 Sentence Compression

🚀 句子壓縮模型

本模型可用於句子壓縮（即抽取式句子摘要）。它會對每個單詞進行預測，判斷該單詞是否可以從句子中刪除而不會嚴重影響句子的含義。生成的句子通常不符合語法規則，但仍然具有一定的實用性。

該模型是基於 rubert-tiny2 在論文 Sentence compression for Russian: dataset and baselines 的數據集上進行微調得到的（數據可在 此處 找到）。

🚀 快速開始

模型功能概述

此模型主要用於俄語句子的壓縮，通過預測每個單詞是否可以被刪除，在一定程度上保留句子的核心含義。

依賴安裝

確保你已經安裝了 torch 和 transformers 庫。如果未安裝，可以使用以下命令進行安裝：

pip install torch transformers

代碼示例

import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer
model_name = 'cointegrated/rubert-tiny2-sentence-compression'
model = AutoModelForTokenClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)


def compress(text, threshold=0.5, keep_ratio=None):
    """ Compress a sentence by removing the least important words.
    Parameters:
        threshold: cutoff for predicted probabilities of word removal
        keep_ratio: proportion of words to preserve
    By default, threshold of 0.5 is used.
    """
    with torch.inference_mode():
        tok = tokenizer(text, return_tensors='pt').to(model.device)
        proba = torch.softmax(model(**tok).logits, -1).cpu().numpy()[0, :, 1]
    if keep_ratio is not None:
        threshold = sorted(proba)[int(len(proba) * keep_ratio)]
    kept_toks = []
    keep = False
    prev_word_id = None
    for word_id, score, token in zip(tok.word_ids(), proba, tok.input_ids[0]):
        if word_id is None:
            keep = True
        elif word_id != prev_word_id:
            keep = score < threshold
        if keep:
            kept_toks.append(token)
        prev_word_id = word_id
    return tokenizer.decode(kept_toks, skip_special_tokens=True)


text = 'Кроме того, можно взять идею, рожденную из сердца, и выразить ее в рамках одной '\
    'из этих структур, без потери искренности идеи и смысла песни.'
    
print(compress(text))
print(compress(text, threshold=0.3))
print(compress(text, threshold=0.1))
# можно взять идею, рожденную из сердца, и выразить ее в рамках одной из этих структур.
# можно взять идею, рожденную из сердца выразить ее в рамках одной из этих структур.
# можно взять идею рожденную выразить структур.

print(compress(text, keep_ratio=0.5))
# можно взять идею, рожденную из сердца выразить ее в рамках структур.

代碼解釋

上述代碼展示瞭如何使用該模型進行句子壓縮。主要步驟如下：

1. 導入必要的庫：導入 torch 和 transformers 庫。
2. 加載模型和分詞器：使用 AutoModelForTokenClassification 和 AutoTokenizer 加載預訓練的模型和分詞器。
3. 定義壓縮函數：定義 compress 函數，該函數接受文本、閾值和保留比例作為參數，返回壓縮後的句子。
4. 測試壓縮函數：定義一個測試文本，並調用 compress 函數進行不同參數的測試。

注意事項

• 生成的句子可能不符合語法規則，但仍然可以保留句子的核心含義。
• 可以根據需要調整 threshold 和 keep_ratio 參數，以獲得不同程度的壓縮效果。