lukecq commited on
Commit
d1884f3
1 Parent(s): 958beef

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +22 -15
README.md CHANGED
@@ -31,8 +31,7 @@ There are three versions of models released. The details are:
31
  |------------|-----------|----------|-------|-------|----|
32
  | [zero-shot-classify-SSTuning-base](https://huggingface.co/DAMO-NLP-SG/zero-shot-classify-SSTuning-base) | [roberta-base](https://huggingface.co/roberta-base) | 125M | Low | High | 20.48M |
33
  | [zero-shot-classify-SSTuning-large](https://huggingface.co/DAMO-NLP-SG/zero-shot-classify-SSTuning-large) | [roberta-large](https://huggingface.co/roberta-large) | 355M | Medium | Medium | 5.12M |
34
- | [zero-shot-classify-SSTuning-ALBERT](https://huggingface.co/DAMO-NLP-SG/zero-shot-classify-SSTuning-ALBERT)
35
- | [albert-xxlarge-v2](https://huggingface.co/albert-xxlarge-v2) | 235M | High | Low| 5.12M |
36
 
37
  Please note that zero-shot-classify-SSTuning-base is trained with more data (20.48M) than the paper, as this will increase the accuracy.
38
 
@@ -49,31 +48,39 @@ You can try the model with the Colab [Notebook](https://colab.research.google.co
49
  from transformers import AutoTokenizer, AutoModelForSequenceClassification
50
  import torch, string, random
51
 
52
- tokenizer = AutoTokenizer.from_pretrained("DAMO-NLP-SG/zero-shot-classify-SSTuning-base")
53
- model = AutoModelForSequenceClassification.from_pretrained("DAMO-NLP-SG/zero-shot-classify-SSTuning-base")
54
 
55
  text = "I love this place! The food is always so fresh and delicious."
56
  list_label = ["negative", "positive"]
57
 
 
58
  list_ABC = [x for x in string.ascii_uppercase]
59
- def add_prefix(text, list_label, shuffle = False):
 
60
  list_label = [x+'.' if x[-1] != '.' else x for x in list_label]
61
  list_label_new = list_label + [tokenizer.pad_token]* (20 - len(list_label))
62
  if shuffle:
63
  random.shuffle(list_label_new)
64
  s_option = ' '.join(['('+list_ABC[i]+') '+list_label_new[i] for i in range(len(list_label_new))])
65
- return f'{s_option} {tokenizer.sep_token} {text}', list_label_new
66
-
67
- text_new, list_label_new = add_prefix(text,list_label,shuffle=False)
68
-
69
- encoding = tokenizer([text_new],truncation=True, padding='max_length',max_length=512, return_tensors='pt')
70
- with torch.no_grad():
71
- logits = model(**encoding).logits
 
72
  probs = torch.nn.functional.softmax(logits, dim = -1).tolist()
73
- predictions = torch.argmax(logits, dim=-1)
 
 
 
 
74
 
75
- print(probs)
76
- print(predictions)
 
77
  ```
78
 
79
 
 
31
  |------------|-----------|----------|-------|-------|----|
32
  | [zero-shot-classify-SSTuning-base](https://huggingface.co/DAMO-NLP-SG/zero-shot-classify-SSTuning-base) | [roberta-base](https://huggingface.co/roberta-base) | 125M | Low | High | 20.48M |
33
  | [zero-shot-classify-SSTuning-large](https://huggingface.co/DAMO-NLP-SG/zero-shot-classify-SSTuning-large) | [roberta-large](https://huggingface.co/roberta-large) | 355M | Medium | Medium | 5.12M |
34
+ | [zero-shot-classify-SSTuning-ALBERT](https://huggingface.co/DAMO-NLP-SG/zero-shot-classify-SSTuning-ALBERT) | [albert-xxlarge-v2](https://huggingface.co/albert-xxlarge-v2) | 235M | High | Low| 5.12M |
 
35
 
36
  Please note that zero-shot-classify-SSTuning-base is trained with more data (20.48M) than the paper, as this will increase the accuracy.
37
 
 
48
  from transformers import AutoTokenizer, AutoModelForSequenceClassification
49
  import torch, string, random
50
 
51
+ tokenizer = AutoTokenizer.from_pretrained("albert-xxlarge-v2")
52
+ model = AutoModelForSequenceClassification.from_pretrained("DAMO-NLP-SG/zero-shot-classify-SSTuning-ALBERT")
53
 
54
  text = "I love this place! The food is always so fresh and delicious."
55
  list_label = ["negative", "positive"]
56
 
57
+ device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
58
  list_ABC = [x for x in string.ascii_uppercase]
59
+
60
+ def check_text(model, text, list_label, shuffle=False):
61
  list_label = [x+'.' if x[-1] != '.' else x for x in list_label]
62
  list_label_new = list_label + [tokenizer.pad_token]* (20 - len(list_label))
63
  if shuffle:
64
  random.shuffle(list_label_new)
65
  s_option = ' '.join(['('+list_ABC[i]+') '+list_label_new[i] for i in range(len(list_label_new))])
66
+ text = f'{s_option} {tokenizer.sep_token} {text}'
67
+
68
+ model.to(device).eval()
69
+ encoding = tokenizer([text],truncation=True, max_length=512,return_tensors='pt')
70
+ item = {key: val.to(device) for key, val in encoding.items()}
71
+ logits = model(**item).logits
72
+
73
+ logits = logits if shuffle else logits[:,0:len(list_label)]
74
  probs = torch.nn.functional.softmax(logits, dim = -1).tolist()
75
+ predictions = torch.argmax(logits, dim=-1).item()
76
+ probabilities = [round(x,5) for x in probs[0]]
77
+
78
+ print(f'prediction: {predictions} => ({list_ABC[predictions]}) {list_label_new[predictions]}')
79
+ print(f'probability: {round(probabilities[predictions]*100,2)}%')
80
 
81
+ check_text(model, text, list_label)
82
+ # prediction: 1 => (B) positive.
83
+ # probability: 98.64%
84
  ```
85
 
86