Update README.md

README.md

@@ -59,20 +59,20 @@ list_label = ["negative", "positive"]
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 list_ABC = [x for x in string.ascii_uppercase]
-def add_prefix(text,list_label, shuffle=False):
+
+def check_text(model, text, list_label, shuffle=False): 
     list_label = [x+'.' if x[-1] != '.' else x for x in list_label]
     list_label_new = list_label + [tokenizer.pad_token]* (20 - len(list_label))
     if shuffle: 
         random.shuffle(list_label_new)
     s_option = ' '.join(['('+list_ABC[i]+') '+list_label_new[i] for i in range(len(list_label_new))])
-    return f'{s_option} {tokenizer.sep_token} {text}', list_label_new
+    text = f'{s_option} {tokenizer.sep_token} {text}'
 
-def check_text(model, text, list_label, shuffle=False): 
-    text, list_label_new = add_prefix(text,list_label, shuffle = shuffle)
     model.to(device).eval()
-    encoding = tokenizer([text],truncation=True, max_length=512)
-    item = {key: torch.tensor(val).to(device) for key, val in encoding.items()}
+    encoding = tokenizer([text],truncation=True, max_length=512,return_tensors='pt')
+    item = {key: val.to(device) for key, val in encoding.items()}
     logits = model(**item).logits
+    
     logits = logits if shuffle else logits[:,0:len(list_label)]
     probs = torch.nn.functional.softmax(logits, dim = -1).tolist()
     predictions = torch.argmax(logits, dim=-1).item()