Disease Named Entity Recognition Romanian (Base, ONNX)

Description

This is an XLM-RoBERTa-based NER model for Romanian medical text. It identifies disease mentions in clinical sentences and labels them using B-DISEASE, I-DISEASE, or O. The model is intended for extracting disease names from notes, reports, and diagnoses, and is provided in ONNX format for efficient inference.

Download Copy S3 URI

How to use

from sparknlp.base import *
from sparknlp_jsl.annotator import *
from pyspark.ml import Pipeline

documentAssembler = DocumentAssembler() \
    .setInputCol("text") \
    .setOutputCol("document")

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

tokenClassifier = XlmRoBertaForTokenClassification \
    .pretrained("xlmroberta_disease_ner_ro_onnx", "ro", "clinical/models") \
    .setInputCols(["document", "token"]) \
    .setOutputCol("ner")

converter = NerConverter() \
    .setInputCols(["document", "token", "ner"]) \
    .setOutputCol("ner_chunk")

pipeline = Pipeline(stages=[
    documentAssembler,
    tokenizer,
    tokenClassifier,
    converter
])

data = spark.createDataFrame([
    ["Pacientul a fost diagnosticat cu pneumonie bacteriană și insuficiență cardiacă, având antecedente de diabet zaharat și hipertensiune arterială."],
    ["Femeia suferă de astm bronșic cronic și boală pulmonară obstructivă, fiind monitorizată pentru insuficiență respiratorie."],
    ["Bărbatul are istoric de ulcer gastric și hepatită virală, iar recent a dezvoltat o infecție respiratorie acută."]
]).toDF("text")

result = pipeline.fit(data).transform(data)

result.selectExpr("explode(ner_chunk) as chunk").selectExpr(
    "chunk.result as text",
    "chunk.metadata['entity'] as entity"
).show(truncate=False)

from johnsnowlabs import nlp, medical

documentAssembler = nlp.DocumentAssembler() \
    .setInputCol("text") \
    .setOutputCol("document")

tokenizer = nlp.Tokenizer() \
    .setInputCols("document") \
    .setOutputCol("token")

tokenClassifier = nlp.XlmRoBertaForTokenClassification \
    .pretrained("xlmroberta_disease_ner_ro_onnx", "ro", "clinical/models") \
    .setInputCols(["document", "token"]) \
    .setOutputCol("ner")

converter = nlp.NerConverter() \
    .setInputCols(["document", "token", "ner"]) \
    .setOutputCol("ner_chunk")

pipeline = nlp.Pipeline(stages=[
    documentAssembler,
    tokenizer,
    tokenClassifier,
    converter
])

data = spark.createDataFrame([
    ["Pacientul a fost diagnosticat cu pneumonie bacteriană și insuficiență cardiacă, având antecedente de diabet zaharat și hipertensiune arterială."],
    ["Femeia suferă de astm bronșic cronic și boală pulmonară obstructivă, fiind monitorizată pentru insuficiență respiratorie."],
    ["Bărbatul are istoric de ulcer gastric și hepatită virală, iar recent a dezvoltat o infecție respiratorie acută."]
]).toDF("text")

result = pipeline.fit(data).transform(data)

result.selectExpr("explode(ner_chunk) as chunk").selectExpr(
    "chunk.result as text",
    "chunk.metadata['entity'] as entity"
).show(truncate=False)

import com.johnsnowlabs.nlp.base._
import com.johnsnowlabs.nlp.annotators._
import org.apache.spark.ml.Pipeline

val documentAssembler = new DocumentAssembler()
  .setInputCol("text")
  .setOutputCol("document")

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

val tokenClassifier = XlmRoBertaForTokenClassification
  .pretrained("xlmroberta_disease_ner_ro_onnx", "ro", "clinical/models")
  .setInputCols(Array("document", "token"))
  .setOutputCol("ner")

val converter = new NerConverter()
  .setInputCols(Array("document", "token", "ner"))
  .setOutputCol("ner_chunk")

val pipeline = new Pipeline().setStages(Array(
  documentAssembler,
  tokenizer,
  tokenClassifier,
  converter
))

val data = spark.createDataFrame(Seq(
  "Pacientul a fost diagnosticat cu pneumonie bacteriană și insuficiență cardiacă, având antecedente de diabet zaharat și hipertensiune arterială.",
  "Femeia suferă de astm bronșic cronic și boală pulmonară obstructivă, fiind monitorizată pentru insuficiență respiratorie.",
  "Bărbatul are istoric de ulcer gastric și hepatită virală, iar recent a dezvoltat o infecție respiratorie acută."
)).toDF("text")

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

result.selectExpr("explode(ner_chunk) as chunk")
  .selectExpr("chunk.result as text", "chunk.metadata['entity'] as entity")
  .show(false)

Results

+-------------------+-------+
|text               |entity |
+-------------------+-------+
|astm bronșic cronic|DISEASE|
|boală pulmonară    |DISEASE|
|insuficiență       |DISEASE|
+-------------------+-------+

Model Information

Model Name:	xlmroberta_disease_ner_ro_onnx
Compatibility:	Healthcare NLP 6.2.0+
License:	Licensed
Edition:	Official
Input Labels:	[token, sentence]
Language:	ro
Size:	1.0 GB

PREVIOUSProcedures Named Entity Recognition Italian (Base, ONNX)

NEXTSymptom Named Entity Recognition Romanian (Base, ONNX)