moshew commited on
Commit
37aab5a
1 Parent(s): b8fa89e

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +33 -1
README.md CHANGED
@@ -11,9 +11,34 @@ pipeline_tag: text-classification
11
 
12
  This is a [SetFit model](https://github.com/huggingface/setfit) that can be used for text classification. The model has been trained using an efficient few-shot learning technique that involves:
13
 
14
- 1. Fine-tuning a [Sentence Transformer](https://www.sbert.net) with contrastive learning.
15
  2. Training a classification head with features from the fine-tuned Sentence Transformer.
16
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
17
  ## Usage
18
 
19
  To use this model for inference, first install the SetFit library:
@@ -33,6 +58,13 @@ model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-englis
33
  preds = model(["i loved the spiderman movie!", "pineapple on pizza is the worst 🤮"])
34
  ```
35
 
 
 
 
 
 
 
 
36
  ## BibTeX entry and citation info
37
 
38
  ```bibtex
 
11
 
12
  This is a [SetFit model](https://github.com/huggingface/setfit) that can be used for text classification. The model has been trained using an efficient few-shot learning technique that involves:
13
 
14
+ 1. Fine-tuning a [Sentence Transformer](https://www.sbert.net) ("BAAI/bge-small-en-v1.5") with contrastive learning.
15
  2. Training a classification head with features from the fine-tuned Sentence Transformer.
16
 
17
+ ## Training code
18
+
19
+ ```python
20
+ from setfit import SetFitModel
21
+
22
+ from datasets import load_dataset
23
+ from setfit import SetFitModel, SetFitTrainer
24
+
25
+ # Load a dataset from the Hugging Face Hub
26
+ dataset = load_dataset("SetFit/sst2")
27
+
28
+ # Upload Train and Test data
29
+ num_classes = 2
30
+ test_ds = dataset["test"]
31
+ train_ds = dataset["train"]
32
+
33
+ model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5")
34
+ trainer = SetFitTrainer(model=model, train_dataset=train_ds, eval_dataset=test_ds)
35
+
36
+ # Train and evaluate
37
+ trainer.train()
38
+ trainer.evaluate()['accuracy']
39
+
40
+ ```
41
+
42
  ## Usage
43
 
44
  To use this model for inference, first install the SetFit library:
 
58
  preds = model(["i loved the spiderman movie!", "pineapple on pizza is the worst 🤮"])
59
  ```
60
 
61
+ ## Accuracy
62
+ On SST-2 dev set:
63
+
64
+ 91.4% SetFit
65
+
66
+ 88.4% (no Fine-Tuning)
67
+
68
  ## BibTeX entry and citation info
69
 
70
  ```bibtex