Update phoBERT.py

phoBERT.py

@@ -2,6 +2,7 @@ import torch
 from transformers import AutoModel, AutoTokenizer
 from underthesea import word_tokenize
 import __main__
+import time
 
 
 #phobert = AutoModel.from_pretrained("vinai/phobert-base")
@@ -69,7 +70,7 @@ def tokenize(data):
   return output
 
 def BERT_predict(text):
-    
+    t1 = time.time()
     text = [text]
     token = tokenize(text)
     
@@ -79,13 +80,15 @@ def BERT_predict(text):
     
     result = model(ids, mask, token_type_ids)
     # print(result)
+    t2 = time.time()
+    print(f'phoBERT: {t2-t1}s')
     return result.tolist()[0]
 
-print(BERT_predict("xin chaof"))
-print(BERT_predict("con chó"))
-print(BERT_predict("đồ chó"))
-print(BERT_predict("đồ ngu"))
-print(BERT_predict("cái lồn"))
-print(BERT_predict("óc chó"))
-print(BERT_predict("đồ chó đẻ"))
-print(BERT_predict("con đĩ"))
+# print(BERT_predict("xin chaof"))
+# print(BERT_predict("con chó"))
+# print(BERT_predict("đồ chó"))
+# print(BERT_predict("đồ ngu"))
+# print(BERT_predict("cái lồn"))
+# print(BERT_predict("óc chó"))
+# print(BERT_predict("đồ chó đẻ"))
+# print(BERT_predict("con đĩ"))