Update README.md

README.md

@@ -25,7 +25,6 @@ This model has been trained on the following datasets:
 from transformers import AutoTokenizer, AutoModelForTokenClassification
 model_id = "gauneg/roberta-base-absa-ate-sentiment"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForTokenClassification.from_pretrained(model_id)
 
 
 # the sequence of labels used during training
@@ -33,6 +32,8 @@ labels = {"B-neu": 1, "I-neu": 2, "O": 0, "B-neg": 3, "B-con": 4, "I-pos": 5, "B
 id2lab = {idx: lab for lab, idx in labels.items()}
 lab2id = {lab: idx for lab, idx in labels.items()}
 
+model = AutoModelForTokenClassification.from_pretrained(model_id, 
+                                                        num_labels=len(labels), id2label=id2lab, label2id=lab2id)
 
 # making one prediction at a time (should be padded/batched and truncated for efficiency)
 text_input = "Been here a few times and food has always been good but service really suffers when it gets crowded."
@@ -41,14 +42,17 @@ tok_inputs = tokenizer(text_input, return_tensors="pt")
 
 y_pred = model(**tok_inputs) # predicting the logits
 
-y_pred_fin = y_pred.logits.argmax(dim=-1)[0] # selecting the most favoured labels for each token from the logits
+# since first and the last tokens are excluded (<s> and </s>)
+# they have to be removed before decoding the labels predicted against them
+y_pred_fin = y_pred.logits.argmax(dim=-1)[0][1:-1] # selecting the most favoured labels for each token from the logits
+
 
 decoded_pred = [id2lab[logx.item()] for logx in y_pred_fin]
 
 
 ## displaying the input tokens with predictions and skipping <s> and </s> tokens at the beginning and the end respectively
-
-tok_levl_pred = list(zip(tokenizer.convert_ids_to_tokens(tok_inputs['input_ids'][0]), decoded_pred))[1:-1]
+decoded_toks = tok_inputs['input_ids'][0][1:-1]
+tok_levl_pred = list(zip(tokenizer.convert_ids_to_tokens(decoded_toks), decoded_pred))
 
 ```