Update README.md
Browse files
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
|