Update Demo.py

Demo.py

@@ -45,9 +45,18 @@ def create_pipeline(model):
         .setInputCols(["document"]) \
         .setOutputCol("token")
 
-    embeddings = WordEmbeddingsModel.pretrained('glove_840B_300', lang='xx') \
-        .setInputCols(["document", "token"]) \
-        .setOutputCol("embeddings")
+    if model == "ner_wikiner_glove_840B_300":
+        embeddings = WordEmbeddingsModel.pretrained('glove_840B_300', lang='xx') \
+            .setInputCols(['document', 'token']) \
+            .setOutputCol('embeddings')
+    elif model == "wikiner_6B_300":
+        embeddings = WordEmbeddingsModel.pretrained('glove_6B_300', lang='xx') \
+            .setInputCols(['document', 'token']) \
+            .setOutputCol('embeddings')
+    elif model == "wikiner_6B_100":
+        embeddings = WordEmbeddingsModel.pretrained('glove_100d') \
+            .setInputCols(['document', 'token']) \
+            .setOutputCol('embeddings')
 
     lang = 'xx' if model == "ner_wikiner_glove_840B_300" else 'es'
     ner_model = NerDLModel.pretrained(model, lang) \