ALBERT Token Classification Base (albert_token_classifier_base_by_hooshvarelab)

Description

ALBERT Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks.

Download Copy S3 URI

How to use



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

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

tokenClassifier = AlbertForTokenClassification .pretrained('albert_token_classifier_base_by_hooshvarelab', 'en') .setInputCols(['token', 'document']) .setOutputCol('ner') .setCaseSensitive(True) .setMaxSentenceLength(512)

# since output column is IOB/IOB2 style, NerConverter can extract entities
ner_converter = NerConverter() .setInputCols(['document', 'token', 'ner']) .setOutputCol('entities')

pipeline = Pipeline(stages=[
document_assembler, 
tokenizer,
tokenClassifier,
ner_converter
])

example = spark.createDataFrame([['My name is John!']]).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 = AlbertForTokenClassification.pretrained("albert_token_classifier_base_by_hooshvarelab", "en")
.setInputCols("document", "token")
.setOutputCol("ner")
.setCaseSensitive(true)
.setMaxSentenceLength(512)

// since output column is IOB/IOB2 style, NerConverter can extract entities
val ner_converter = NerConverter() 
.setInputCols("document", "token", "ner") 
.setOutputCol("entities")

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

val example = Seq.empty["My name is John!"].toDS.toDF("text")

val result = pipeline.fit(example).transform(example)

Model Information

Model Name: albert_token_classifier_base_by_hooshvarelab
Compatibility: Spark NLP 5.4.2+
License: Open Source
Edition: Official
Input Labels: [token, document]
Output Labels: [label]
Language: en
Size: 41.9 MB
Case sensitive: true