support new transformers pipeline version

README.md

@@ -33,14 +33,13 @@ Person, Location, Organization, Nationality, Job, Product, Event, Time, Art-Work
 
 Install the following Python packages
 
-`$ pip3 install simpletransformers==0.61.5 nltk==3.5 protobuf==3.15.3 torch==1.7.1`
+`$ pip3 install transformers==4.7.0 nltk==3.5 protobuf==3.15.3 torch==1.7.1`
 
 > If you are using `Google Colab`, please restart your runtime after installing the packages.
 
 -----------
 
 ```python
-from simpletransformers.ner import NERModel, NERArgs
 import logging
 import re
 
@@ -50,89 +49,67 @@ from nltk.tokenize import word_tokenize
 
 # disable INFO Logs
 transformers_logger = logging.getLogger("transformers")
-simpletransformers_logger = logging.getLogger("simpletransformers")
-simpletransformers_ner_logger = logging.getLogger("simpletransformers.ner")
 transformers_logger.setLevel(logging.WARNING)
-simpletransformers_logger.setLevel(logging.WARNING)
-simpletransformers_ner_logger.setLevel(logging.WARNING)
 
-# Load the Model
 custom_labels = ["O", "B-job", "I-job", "B-nationality", "B-person", "I-person", "B-location",
                  "B-time", "I-time", "B-event", "I-event", "B-organization", "I-organization",
                  "I-location", "I-nationality", "B-product", "I-product", "B-artwork", "I-artwork"]
 
-model_args = NERArgs()
-model_args.labels_list=custom_labels
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+from transformers import pipeline
 
-ner_model = NERModel(
-     "xlmroberta", "marefa-nlp/marefa-ner",
-     args=model_args,
-     use_cuda=True # set to False to use CPU
-)
+# ===== import the model
+m_name = "marefa-nlp/marefa-ner"
+tokenizer = AutoTokenizer.from_pretrained(m_name)
+model = AutoModelForTokenClassification.from_pretrained(m_name)
+
+ar_ner = pipeline("ner", model=model, tokenizer=tokenizer, grouped_entities=True, aggregation_strategy="simple")
 
 # Model Inference
 samples = [
     "تلقى تعليمه في الكتاب ثم انضم الى الأزهر عام 1873م. تعلم على يد السيد جمال الدين الأفغاني والشيخ محمد عبده",
     "بعد عودته إلى القاهرة، التحق نجيب الريحاني فرقة جورج أبيض، الذي كان قد ضمَّ - قُبيل ذلك - فرقته إلى فرقة سلامة حجازي . و منها ذاع صيته",
     "امبارح اتفرجت على مباراة مانشستر يونايتد مع ريال مدريد في غياب الدون كرستيانو رونالدو",
-    "Government extends flight ban from India, Pakistan until June 21"
+    "Government extends flight ban from India and Pakistan until June 21"
 ]
 
 # Preprocess
 samples = [ " ".join(word_tokenize(sample.strip())) for sample in samples if sample.strip() != "" ]
 
-# Predict
-predictions, raw_outputs = ner_model.predict(samples)
-
-# Group the Predicted Entities
-entities = []
-for pred in predictions:
-  grouped_entities = []
-  for rec in pred:
-
-    token = list(rec.keys())[0]
-    label = rec[token]
-
-    if label == "O":
-      continue
-    
-    if "B-" in label:
-      grouped_entities.append({"token": token, "label": label.replace("B-","")})
-    elif "I-" in label and len(grouped_entities) > 0:
-      grouped_entities[-1]["token"] += f" {token}"
-  
-  entities.append(grouped_entities)
-  
-# Print the model outputs
-for sample, results in zip(samples, entities):
-  print(sample)
-  for res in results:
-    print("\t", res["token"], "=>", res["label"])
-  print("==================")
+for sample in samples:
+    results = ar_ner(sample)
+    print(sample)
+    for result in results:
+        print("\t", result["word"], "=>", result["entity_group"])
+        
+    print("=========\n")
   
 ###
 # تلقى تعليمه في الكتاب ثم انضم الى الأزهر عام 1873م . تعلم على يد السيد جمال الدين الأفغاني والشيخ محمد عبده
-#  الأزهر => organization
-#  عام 1873م => time
-#  جمال الدين الأفغاني => person
-#  محمد عبده => person
-# ==================
+# 	 الأزهر => organization
+# 	 عام 1873م => time
+# 	 جمال الدين الأفغاني => person
+# 	 محمد عبده => person
+# =========
+
 # بعد عودته إلى القاهرة، التحق نجيب الريحاني فرقة جورج أبيض، الذي كان قد ضمَّ - قُبيل ذلك - فرقته إلى فرقة سلامة حجازي . و منها ذاع صيته
-#  القاهرة، => location
-#  نجيب الريحاني => person
-#  فرقة جورج أبيض، => organization
-#  فرقة سلامة حجازي => organization
-# ==================
+# 	 القاهرة => location
+# 	 نجيب الريحاني => person
+# 	 فرقة جورج أبيض => organization
+# 	 فرقة سلامة حجازي => organization
+# =========
+
 # امبارح اتفرجت على مباراة مانشستر يونايتد مع ريال مدريد في ��ياب الدون كرستيانو رونالدو
-#  مانشستر يونايتد => organization
-#  ريال مدريد => organization
-#  كرستيانو رونالدو => person
-# ==================
-# Government extends flight ban from India , Pakistan until June 21
-#  India => location
-#  Pakistan => location
-#  June 21 => time
-# ==================
+# 	 مانشستر يونايتد => organization
+# 	 ريال مدريد => organization
+# 	 كرستيانو رونالدو => person
+# =========
+
+# Government extends flight ban from India and Pakistan until June 21
+# 	 India => location
+# 	 Pakistan => location
+# 	 June 21 => time
+# =========
 ###
 ```