Update app.py

app.py

@@ -61,18 +61,43 @@ st.sidebar.markdown(logo_html, unsafe_allow_html=True)
 
 
 #sidebar info
-model_name= ["nerdl_fewnerd_100d"]
+model_name= ["nerdl_fewnerd_100d", "bert_large_token_classifier_ontonote", "ner_mit_movie_complex_distilbert_base_cased", "ner_conll_albert_large_uncased"]
 st.sidebar.title("Pretrained model to test")
 selected_model = st.sidebar.selectbox("", model_name)
 
 ######## Main Page #########
-app_title= "Detect up to 8 entity types in general domain texts"
-app_description= "Named Entity Recognition model aimed to detect up to 8 entity types from general domain texts. This model was trained on the Few-NERD/inter public dataset using Spark NLP, and is available in Spark NLP Models hub (https://nlp.johnsnowlabs.com/models)"
-st.title(app_title)
-st.markdown("<h2>"+app_description+"</h2>" , unsafe_allow_html=True)
 if selected_model == "nerdl_fewnerd_100d":
+    app_title= "Detect up to 8 entity types in general domain texts"
+    app_description= "Named Entity Recognition model aimed to detect up to 8 entity types from general domain texts. This model was trained on the Few-NERD/inter public dataset using Spark NLP, and it is available in Spark NLP Models hub (https://nlp.johnsnowlabs.com/models)"
+    st.title(app_title)
+    st.markdown("<h2>"+app_description+"</h2>" , unsafe_allow_html=True)
     st.markdown("**`PERSON`**   **,** **`ORGANIZATION`**    **,**  **`LOCATION`** **,**  **`ART`** **,**  **`BUILDING`** **,**  **`PRODUCT`** **,**  **`EVENT`** **,**  **`OTHER`**", unsafe_allow_html=True)
 
+elif selected_model== "bert_large_token_classifier_ontonote":
+    app_title= "Detect up to 18 entity types in general domain texts"
+    app_description= "Named Entity Recognition model aimed to detect up to 18 entity types from general domain texts. This model is a fine-tuned BERT model that is ready to use for Named Entity Recognition and achieves state-of-the-art performance for the NER task, and it is available in Spark NLP Models hub (https://nlp.johnsnowlabs.com/models)"
+    st.title(app_title)
+    st.markdown("<h2>"+app_description+"</h2>" , unsafe_allow_html=True)
+    st.markdown("""**`CARDINAL`**   **,** **`DATE`**    **,**  **`EVENT`** **,**  **`FAC`** **,**  **`GPE`** **,**  **`LANGUAGE`** **,**  **`LAW`** **,**  **`LOC`**,
+                **`MONEY`**   **,** **`NORP`**    **,**  **`ORDINAL`** **,**  **`ORG`** **,**  **`PERCENT`** **,**  **`PERCENT`** **,**  **`PERSON`** **,**  **`PRODUCT`**,
+                **`QUANTITY`**   **,** **`TIME`**    **,**  **`WORK_OF_ART` **""", unsafe_allow_html=True)
+
+
+elif selected_model== "ner_mit_movie_complex_distilbert_base_cased":
+    app_title= "Detect up to 12 entity types in movie domain texts"
+    app_description= "Named Entity Recognition model aimed to detect up to 12 entity types from movie domain texts. This model was trained on the MIT Movie Corpus complex queries dataset to detect movie trivia using Spark NLP, and it is available in Spark NLP Models hub (https://nlp.johnsnowlabs.com/models)"
+    st.title(app_title)
+    st.markdown("<h2>"+app_description+"</h2>" , unsafe_allow_html=True)
+    st.markdown("""**`ACTOR`**   **,** **`AWARD`**    **,**  **`CHARACTER_NAME`** **,**  **`DIRECTOR`** **,**  **`GENRE`** **,**  **`OPINION`** **,**  **`ORIGIN`** **,**  **`PLOT`**,
+                    **`QUOTE`**   **,** **`RELATIONSHIP`**    **,**  **`SOUNDTRACK`** **,**  **`YEAR`  **""", unsafe_allow_html=True)
+
+elif selected_model=="ner_conll_albert_large_uncased":
+    app_title= "Detect up to 4 entity types in general domain texts"
+    app_description= "Named Entity Recognition model aimed to detect up to 4 entity types from general domain texts. This model was trained on the CoNLL 2003 text corpus using Spark NLP, and it is available in Spark NLP Models hub (https://nlp.johnsnowlabs.com/models)"
+    st.title(app_title)
+    st.markdown("<h2>"+app_description+"</h2>" , unsafe_allow_html=True)
+    st.markdown("**`PER`**   **,** **`LOC`**    **,**  **`ORG`** **,**  **`MISC` **", unsafe_allow_html=True)
+
 st.subheader("")
 
 
@@ -91,31 +116,103 @@ def get_pipeline(text):
         .setInputCols(["sentence"])\
         .setOutputCol("token")
 
-    embeddings= WordEmbeddingsModel.pretrained("glove_100d")\
-        .setInputCols(["sentence", "token"])\
-        .setOutputCol("embeddings")
-
-
-    ner= NerDLModel.pretrained("nerdl_fewnerd_100d")\
-        .setInputCols(["document", "token", "embeddings"])\
-        .setOutputCol("ner")
 
+    if selected_model=="nerdl_fewnerd_100d":
+        embeddings= WordEmbeddingsModel.pretrained("glove_100d")\
+            .setInputCols(["sentence", "token"])\
+            .setOutputCol("embeddings")
+
+        ner= NerDLModel.pretrained(selected_model)\
+            .setInputCols(["document", "token", "embeddings"])\
+            .setOutputCol("ner")
+
+
+        ner_converter= NerConverter()\
+            .setInputCols(["sentence", "token", "ner"])\
+            .setOutputCol("ner_chunk")
+
+
+        pipeline = Pipeline(
+            stages = [
+                documentAssembler,
+                sentenceDetector,
+                tokenizer,
+                embeddings,
+                ner,
+                ner_converter
+            ])
+
+    elif selected_model=="bert_large_token_classifier_ontonote":
+        tokenClassifier = BertForTokenClassification \
+            .pretrained('bert_large_token_classifier_ontonote', 'en') \
+            .setInputCols(['token', 'document']) \
+            .setOutputCol('ner') \
+            .setCaseSensitive(True) \
+            .setMaxSentenceLength(512)
+
+        ner_converter= NerConverter()\
+            .setInputCols(["document", "token", "ner"])\
+            .setOutputCol("ner_chunk")
+        
+        pipeline = Pipeline(
+            stages = [
+                documentAssembler,
+                sentenceDetector,
+                tokenizer,
+                tokenClassifier,
+                ner_converter
+            ])
+
+    elif selected_model=="ner_mit_movie_complex_distilbert_base_cased":
+        embeddings = DistilBertEmbeddings\
+            .pretrained('distilbert_base_cased', 'en')\
+            .setInputCols(["token", "document"])\
+            .setOutputCol("embeddings")
+
+        ner = NerDLModel.pretrained('ner_mit_movie_complex_distilbert_base_cased', 'en') \
+            .setInputCols(['document', 'token', 'embeddings']) \
+            .setOutputCol('ner')
+
+        ner_converter= NerConverter()\
+            .setInputCols(["document", "token", "ner"])\
+            .setOutputCol("ner_chunk")
+        
+        pipeline = Pipeline(
+            stages = [
+                documentAssembler,
+                sentenceDetector,
+                tokenizer,
+                embeddings,
+                ner,
+                ner_converter
+            ])
+
+    elif selected_model=="ner_conll_albert_large_uncased":
+        embeddings = AlbertEmbeddings\
+            .pretrained('albert_large_uncased', 'en')\
+            .setInputCols(["document", "token"])\
+            .setOutputCol("embeddings")
+
+        ner = NerDLModel.pretrained('ner_conll_albert_large_uncased', 'en') \
+            .setInputCols(['document', 'token', 'embeddings']) \
+            .setOutputCol('ner')
+
+        ner_converter = NerConverter()\
+            .setInputCols(["document","token","ner"])\
+            .setOutputCol("ner_chunk")
+
+        pipeline = Pipeline(
+            stages = [
+                documentAssembler,
+                sentenceDetector,
+                tokenizer,
+                embeddings,
+                ner,
+                ner_converter
+            ])
 
-    ner_converter= NerConverter()\
-        .setInputCols(["sentence", "token", "ner"])\
-        .setOutputCol("ner_chunk")
 
 
-    pipeline = Pipeline(
-        stages = [
-        documentAssembler,
-        sentenceDetector,
-        tokenizer,
-        embeddings,
-        ner,
-        ner_converter
-        ])
-
     empty_df = spark.createDataFrame([[""]]).toDF("text")
     pipeline_model = pipeline.fit(empty_df)
 
@@ -126,7 +223,13 @@ def get_pipeline(text):
 
 
     
-text= st.text_input("Type here your text and press enter to run:", value="12 Corazones ('12 Hearts') is Spanish-language dating game show produced in the United States for the television network Telemundo since January 2005, based on its namesake Argentine TV show format. The show is filmed in Los Angeles and revolves around the twelve Zodiac signs that identify each contestant. In 2008, Ho filmed a cameo in the Steven Spielberg feature film The Cloverfield Paradox, as a news pundit.")
+if selected_model=="ner_mit_movie_complex_distilbert_base_cased":
+    text= st.text_input("Type here your text and press enter to run:", value="It's only appropriate that Solaris, Russian filmmaker Andrei Tarkovsky's psychological sci-fi classic from 1972, contains an equally original and mind-bending score. Solaris explores the inadequacies of time and memory on an enigmatic planet below a derelict space station. To reinforce the film's chilling setting, Tarkovsky commissioned composer Eduard Artemiev to construct an electronic soundscape reflecting planet Solaris' amorphous and mysterious surface")
+
+
+else:
+    text= st.text_input("Type here your text and press enter to run:", value="12 Corazones ('12 Hearts') is Spanish-language dating game show produced in the United States for the television network Telemundo since January 2005, based on its namesake Argentine TV show format. The show is filmed in Los Angeles and revolves around the twelve Zodiac signs that identify each contestant. In 2008, Ho filmed a cameo in the Steven Spielberg feature film The Cloverfield Paradox, as a news pundit.")
+
 
 #placeholder for warning
 placeholder= st.empty()
@@ -151,3 +254,9 @@ labels = st.sidebar.multiselect(
 
 show_html2(text, df, labels, "Text annotated with identified Named Entities")
 
+
+try_link="""<a href="https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/Certification_Trainings/Public/3.SparkNLP_Pretrained_Models.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/></a>"""
+st.sidebar.title('')
+st.sidebar.markdown('Try it yourself:')
+st.sidebar.markdown(try_link, unsafe_allow_html=True)
+