roberthsu2003 commited on
Commit
7d467f2
·
verified ·
1 Parent(s): 58ff653

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +3 -1
README.md CHANGED
@@ -71,7 +71,9 @@ model = AutoModelForTokenClassification.from_pretrained('roberthsu2003/models_fo
71
  tokenizer = AutoTokenizer.from_pretrained('roberthsu2003/models_for_ner')
72
 
73
  # The label mapping (you might need to adjust this based on your training)
74
- label_list = ['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC']
 
 
75
 
76
  def predict_ner(text):
77
  """Predicts NER tags for a given text using the loaded model."""
 
71
  tokenizer = AutoTokenizer.from_pretrained('roberthsu2003/models_for_ner')
72
 
73
  # The label mapping (you might need to adjust this based on your training)
74
+ #['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC']
75
+ label_list = list(model.config.id2label.values())
76
+
77
 
78
  def predict_ner(text):
79
  """Predicts NER tags for a given text using the loaded model."""