Update README.md

README.md

@@ -56,17 +56,18 @@ inputs = tokenizer(tokens, is_split_into_words=True, return_tensors="pt", trunca
 outputs = model(**inputs)
 predictions = torch.argmax(outputs.logits, dim=2)
 
-# Load tag mappings
-import json
-with open("id2tag.json", "r") as f:
-    id2tag = json.load(f)
+# Get tag mapping from model config
+id2tag = model.config.id2label
 
-# Convert predictions to tags
 word_ids = inputs.word_ids()
 prev_word = None
 for idx, word_idx in enumerate(word_ids):
     if word_idx is not None and word_idx != prev_word:
-        tag = id2tag[str(predictions[0][idx].item())]
+        tag_id = predictions[0][idx].item()
+        if isinstance(id2tag, dict):
+            tag = id2tag.get(str(tag_id), id2tag.get(tag_id, "UNK"))
+        else:
+            tag = id2tag[tag_id] if tag_id < len(id2tag) else "UNK"
         print(tokens[word_idx], "->", tag)
     prev_word = word_idx
 ```