simplified example code

README.md

@@ -17,20 +17,17 @@ tokenizer = AutoTokenizer.from_pretrained("lumalik/vent-roberta-emotion")
 model = AutoModelForSequenceClassification.from_pretrained("lumalik/vent-roberta-emotion")
 model.eval()
 
-texts = ["I love her sooo much", "I hate you!"]
+texts = ["You wont believe what happened to me today",
+         "You wont believe what happened to me today!",
+         "You wont believe what happened to me today...",
+         "You wont believe what happened to me today <3",
+         "You wont believe what happened to me today :)",
+         "You wont believe what happened to me today :("]
 
 for text in texts:
-    encoded_text = tokenizer.encode_plus(text,
-                                         add_special_tokens=True,
-                                         max_length=128,
-                                         return_token_type_ids=True,
-                                         padding="max_length",
-                                         truncation=True,
-                                         return_attention_mask=True)
-
-    output = model(input_ids=torch.tensor(encoded_text['input_ids'], dtype=torch.long).unsqueeze(0),
-                   token_type_ids=torch.tensor(encoded_text['token_type_ids'], dtype=torch.long).unsqueeze(0),
-                   attention_mask=torch.tensor(encoded_text['attention_mask'], dtype=torch.long).unsqueeze(0))
+    encoded_text = tokenizer(text, return_tensors="pt")
+
+    output = model(**encoded_text)
 
     output = softmax(output[0].detach().numpy(), axis=1)