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 参数，以获得不同程度的压缩效果。