Update Demo.py

Demo.py

@@ -42,31 +42,27 @@ def init_spark():
 
 @st.cache_resource
 def create_pipeline(model):
-    document_assembler = DocumentAssembler() \
-        .setInputCol("text") \
-        .setOutputCol("document")
+    documentAssembler = DocumentAssembler()\
+    .setInputCol("text")\
+    .setOutputCol("document")
 
-    sentence_detector = SentenceDetector() \
-        .setInputCols(["document"]) \
-        .setOutputCol("sentence")
+    sentenceDetector = SentenceDetectorDLModel.pretrained("sentence_detector_dl", "xx")\
+    .setInputCols(["document"])\
+    .setOutputCol("sentence")
 
-    word_segmenter = WordSegmenterModel.pretrained("wordseg_large", "zh") \
+    tokenizer = WordSegmenterModel.pretrained("wordseg_large", "zh") \
         .setInputCols(["sentence"]) \
         .setOutputCol("token")
 
-    embeddings = BertEmbeddings.pretrained(name='bert_base_chinese', lang='zh') \
-        .setInputCols(["document", "token"]) \
-        .setOutputCol("embeddings")
+    tokenClassifier = XlmRoBertaForTokenClassification.pretrained("xlm_roberta_large_token_classifier_hrl", "xx")\
+    .setInputCols(["sentence",'token'])\
+    .setOutputCol("ner")
 
-    ner = NerDLModel.pretrained(model, "zh") \
-        .setInputCols(["document", "token", "embeddings"]) \
-        .setOutputCol("ner")
+    ner_converter = NerConverter()\
+    .setInputCols(["sentence", "token", "ner"])\
+    .setOutputCol("ner_chunk")
 
-    ner_converter = NerConverter() \
-        .setInputCols(["sentence", "token", "ner"]) \
-        .setOutputCol("entities")
-
-    pipeline = Pipeline(stages=[document_assembler, sentence_detector, word_segmenter, embeddings, ner, ner_converter])
+    nlpPipeline = Pipeline(stages=[documentAssembler, sentenceDetector, tokenizer, tokenClassifier, ner_converter])
     return pipeline
 
 def fit_data(pipeline, data):
@@ -91,16 +87,11 @@ def annotate(data):
 
 # Set up the page layout
 st.markdown('<div class="main-title">Recognize entities in Chinese text</div>', unsafe_allow_html=True)
-st.markdown("""
-<div class="section">
-    <p>This demo utilizes embeddings-based NER model for Urdu texts, using the urduvec_140M_300d word embeddings</p>
-</div>
-""", unsafe_allow_html=True)
 
 # Sidebar content
 model = st.sidebar.selectbox(
     "Choose the pretrained model",
-    ["ner_msra_bert_768d", "ner_weibo_bert_768d"],
+    ["xlm_roberta_large_token_classifier_hrl"],
     help="For more info about the models visit: https://sparknlp.org/models"
 )