Update README.md

README.md

@@ -25,7 +25,7 @@ The model can be used for zero-shot text classification such sentiment analysis
 The number of labels should be 2 ~ 20. 
 
 ### How to use
-You can try the model with the colab [notebook](https://colab.research.google.com/drive/17bqc8cXFF-wDmZ0o8j7sbrQB9Cq7Gowr?usp=sharing).
+You can try the model with the Colab [Notebook](https://colab.research.google.com/drive/17bqc8cXFF-wDmZ0o8j7sbrQB9Cq7Gowr?usp=sharing).
 
 ```python
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
@@ -35,12 +35,12 @@ tokenizer = AutoTokenizer.from_pretrained("DAMO-NLP-SG/zero-shot-classify-SSTuni
 model = AutoModelForSequenceClassification.from_pretrained("DAMO-NLP-SG/zero-shot-classify-SSTuning-base")
 
 text = "I love this place! The food is always so fresh and delicious."
-list_label = ["negative","positve"]
+list_label = ["negative", "positive"]
 
 list_ABC = [x for x in string.ascii_uppercase]
-def add_prefix(text, list_label, label_num = 20, 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]* (label_num - len(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))])
@@ -50,9 +50,12 @@ 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(**item).logits
+    logits = model(**encoding).logits
     probs = torch.nn.functional.softmax(logits, dim = -1).tolist()
     predictions = torch.argmax(logits, dim=-1)
+
+print(probs)
+print(predictions)
 ```