如何在Python中使用BERT进行文本分类

发布于08月02日

我正在研究一个文本分类问题，并希望使用BERT模型来改进我的结果.我读过有关BERT模型的文章，但我不确定如何在我的特定用例中用Python实现它.

以下是我目前在一个简单的逻辑回归模型中使用的代码:

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline

text_clf = Pipeline([
('tfidf', TfidfVectorizer()),
('clf', LogisticRegression()),
])

text_clf.fit(X_train, y_train)
predictions = text_clf.predict(X_test)

有没有人能举个例子，说明我如何用文本分类的BERT模型来代替它？

推荐答案

嗨，我的朋友，欢迎来到Stackoverflow使用Python中的转换器库，这里是一个基本的例子，告诉你如何使用BERT进行文本分类

从转换器导入BertTokenizer、BertForSequenceClass 进口 flashlight

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
input_ids = tokenizer.encode("Here is some text to classify", add_special_tokens=True)
input_ids = torch.tensor(input_ids).unsqueeze(0)  # Batch size 1
outputs = model(input_ids)
_, predicted_class = torch.max(outputs.logits, dim=1)

print(predicted_class)

查阅transformers个图书馆文档