Update README.md

README.md

@@ -65,10 +65,10 @@ from seqeval.metrics.sequence_labeling import get_entities
 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
 
 # Load model from HuggingFace Hub
-tokenizer = AutoTokenizer.from_pretrained("shibing624/bert4ner-base-uncased")
-model = AutoModelForTokenClassification.from_pretrained("shibing624/bert4ner-base-uncased")
+tokenizer = AutoTokenizer.from_pretrained("../bert4ner-base-uncased")
+model = AutoModelForTokenClassification.from_pretrained("../bert4ner-base-uncased")
 label_list = ["E-ORG", "E-LOC", "S-MISC", "I-MISC", "S-PER", "E-PER", "B-MISC", "O", "S-LOC",
-                "E-MISC", "B-ORG", "S-ORG", "I-ORG", "B-LOC", "I-LOC", "B-PER", "I-PER"]
+              "E-MISC", "B-ORG", "S-ORG", "I-ORG", "B-LOC", "I-LOC", "B-PER", "I-PER"]
 
 sentence = "AL-AIN, United Arab Emirates 1996-12-06"
 
@@ -79,15 +79,15 @@ def get_entity(sentence):
     with torch.no_grad():
         outputs = model(inputs).logits
     predictions = torch.argmax(outputs, dim=2)
-    char_tags = [(token, label_list[prediction]) for token, prediction in zip(tokens, predictions[0].numpy())][1:-1]
+    word_tags = [(token, label_list[prediction]) for token, prediction in zip(tokens, predictions[0].numpy()[1:-1])]
     print(sentence)
-    print(char_tags)
+    print(word_tags)
 
-    pred_labels = [i[1] for i in char_tags]
+    pred_labels = [i[1] for i in word_tags]
     entities = []
     line_entities = get_entities(pred_labels)
     for i in line_entities:
-        word = sentence[i[1]: i[2] + 1]
+        word = tokens[i[1]: i[2] + 1]
         entity_type = i[0]
         entities.append((word, entity_type))