【python transformer】使用Transformer模型进行文本分类

Happy Transformer是基于构建的软件包，可轻松利用最新的NLP模型。高效简单使用，简单训练。下面我们一起来学习学习。

步骤一：安装Happy Transformer

pip install happytransformer
1

步骤二：测试例子

# -*- coding: utf-8 -*-
from happytransformer import HappyTextClassification
# 加载预训练模型
happy_tc = HappyTextClassification("BERT", "ProsusAI/finbert", num_labels=3)
result = happy_tc.classify_text("Tesla's stock just increased by 20%")
print(result)

print(result.label)
print(result.score)
1
2
3
4
5
6
7
8
9

运行结果如下：

TextClassificationResult(label='positive', score=0.9291105270385742)
positive
0.9291105270385742
1
2
3

既可以输出标签又可以输出概率。

步骤四：下面来学习下如何训练文本分类模型

我们必须创建一个包含两列的CSV文件：文本和标签。文本列包含我们希望分类的文本。“标签”列将标签类型包含为大于或等于0的整数。下表给出了训练CSV的示例。
|text| label|
|Tesla's stock just increased by 20%|0|
|The price of gold just dropped by 5%  | 1 |

1
2
3
4
5

训练文本分类算法：

# -*- coding: utf-8 -*-
from happytransformer import HappyTextClassification
from happytransformer import TCTrainArgs
happy_tc = HappyTextClassification(model_type="DISTILBERT", model_name="distilbert-base-uncased", num_labels=2)
args = TCTrainArgs(num_train_epochs=5)
happy_tc.train("train.csv",args=args)
# 预测
result = happy_tc.classify_text("Tesla's stock just increased by 20%")
print(result)
print(result.label)
print(result.score)
1
2
3
4
5
6
7
8
9
10
11

就这样！你刚刚学习了如何实现和训练文本分类Transformer模型。使用Happy Transformer只需几行就可以完成这么多工作!。