Update README.md

README.md

@@ -10,7 +10,8 @@ library_name: transformers
 ---
 
 ### Model summary
-`bert-sci-am` is a BERT-family model trained for scientific literature argument mining. At low-level it performs sequence classification.
+
+`bert-sci-am` is a BERT-family model trained for scientific literature argument mining. At low-level it performs sequence classification. This version is trained on 3-class classification on (david-inf/am-nlp-abstrct)[david-inf/am-nlp-abstrct] forked from [pie/abstrct](https://huggingface.co/datasets/pie/abstrct/tree/main) dataset.
 
 ### How to use
 
@@ -21,7 +22,7 @@ def load_model():
     """Load model from hub"""
     checkpoint = "david-inf/bert-sci-am"
     model = AutoModelForSequenceClassification.from_pretrained(
-        checkpoint, num_labels=2)
+        checkpoint, num_labels=3)
     tokenizer = AutoTokenizer.from_pretrained(checkpoint)
     return model, tokenizer