BERT Sequence Classification - Classify into News Categories

Description

This model is imported from Hugging Face-models. It is a BERT-Mini fine-tuned version of the age_news dataset.

Predicted Entities

World, Sports, Business, Sci/Tech

Download Copy S3 URI

How to use

document_assembler = DocumentAssembler() \
    .setInputCol('text') \
    .setOutputCol('document')

tokenizer = Tokenizer() \
    .setInputCols(['document']) \
    .setOutputCol('token')

sequenceClassifier = BertForSequenceClassification \
      .pretrained('bert_sequence_classifier_age_news', 'en') \
      .setInputCols(['token', 'document']) \
      .setOutputCol('class') \
      .setCaseSensitive(True) \
      .setMaxSentenceLength(512)

pipeline = Pipeline(stages=[document_assembler, tokenizer, sequenceClassifier])

example = spark.createDataFrame([['Microsoft has taken its first step into the metaverse.']]).toDF("text")
result = pipeline.fit(example).transform(example)
val document_assembler = DocumentAssembler() 
    .setInputCol("text") 
    .setOutputCol("document")

val tokenizer = Tokenizer() 
    .setInputCols("document") 
    .setOutputCol("token")

val tokenClassifier = BertForSequenceClassification.pretrained("bert_sequence_classifier_age_news", "en")
      .setInputCols("document", "token")
      .setOutputCol("class")
      .setCaseSensitive(true)
      .setMaxSentenceLength(512)

val pipeline = new Pipeline().setStages(Array(document_assembler, tokenizer, sequenceClassifier))

val example = Seq.empty["Microsoft has taken its first step into the metaverse."].toDS.toDF("text")

val result = pipeline.fit(example).transform(example)
import nlu
nlu.load("en.classify.bert.news.").predict("""Microsoft has taken its first step into the metaverse.""")

Results

['Sci/Tech']

Model Information

Model Name: bert_sequence_classifier_age_news
Compatibility: Spark NLP 3.3.2+
License: Open Source
Edition: Official
Input Labels: [token, sentence]
Output Labels: [label]
Language: en
Case sensitive: true

Data Source

https://huggingface.co/mrm8488/bert-mini-finetuned-age_news-classification

Benchmarking

Test set accuracy: 0.93