Update README.md

README.md

@@ -19,6 +19,7 @@ The model backbone is RoBERTa-base.
 
 The model is tuned with unlabeled data using a learning objective called first sentence prediction (FSP). 
 The FSP task is designed by considering both the nature of the unlabeled corpus and the input/output format of classification tasks. 
+
 The training and validation sets are constructed from the unlabeled corpus using FSP. 
 
 During tuning, BERT-like pre-trained masked language 
@@ -56,8 +57,9 @@ model = AutoModelForSequenceClassification.from_pretrained("DAMO-NLP-SG/zero-sho
 text = "I love this place! The food is always so fresh and delicious."
 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 add_prefix(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: 
@@ -65,16 +67,23 @@ def add_prefix(text, list_label, shuffle = False):
     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_new, list_label_new = add_prefix(text,list_label,shuffle=False)
-
-encoding = tokenizer([text_new],truncation=True, padding='max_length',max_length=512, return_tensors='pt')
-with torch.no_grad():
-    logits = model(**encoding).logits
+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()}
+    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)
+    predictions = torch.argmax(logits, dim=-1).item() 
+    probabilities = [round(x,5) for x in probs[0]]
+
+    print(f'prediction:    {predictions} => ({list_ABC[predictions]}) {list_label_new[predictions]}')
+    print(f'probability:   {round(probabilities[predictions]*100,2)}%')
 
-print(probs)
-print(predictions)
+check_text(model, text, list_label)
+# prediction:    1 => (B) positive.
+# probability:   99.92%
 ```
 
 
@@ -89,8 +98,8 @@ print(predictions)
                Chip Hong Chang and 
                Lidong Bing},
   title     = {Zero-Shot Text Classification via Self-Supervised Tuning},
-  booktitle = {Findings of the 2023 ACL},
+  booktitle = {Findings of the Association for Computational Linguistics: ACL 2023},
   year      = {2023},
-  url       = {},
+  url       = {https://arxiv.org/abs/2305.11442},
 }
 ```