Update README.md

README.md

@@ -11,34 +11,67 @@ model-index:
   results: []
 ---
 
-<!-- This model card has been generated automatically according to the information the Trainer had access to. You
-should probably proofread and complete it, then remove this comment. -->
 
 # stack-edu-classifier-ruby
+This is a classifier for scoring the educational value of code files in The Stack v2 dataset, it is a finetuned version of [bigcode/starencoder](https://huggingface.co/bigcode/starencoder) with a classification head on code files annotated by Llama3.1-70B-Instruct. We use this classifier for building Stack-Edu dataset used for training SmolLM2, see [paper](https://arxiv.org/pdf/2502.02737). Each classifier is trained on one programming language.
 
-This model is a fine-tuned version of [bigcode/starencoder](https://huggingface.co/bigcode/starencoder) on an unknown dataset.
-It achieves the following results on the evaluation set:
-- Loss: 0.3282
-- Precision: 0.4623
-- Recall: 0.3260
-- F1 Macro: 0.3536
-- Accuracy: 0.6657
-- F1 Binary Minimum3: 0.6101
+### How to use in transformers
+To load the classifier, use the following code:
+
+```python
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+tokenizer = AutoTokenizer.from_pretrained(REPO_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(REPO_NAME)
 
-## Model description
+text = "This is a test sentence."
+inputs = tokenizer(text, return_tensors="pt", padding="longest", truncation=True)
+outputs = model(**inputs)
+logits = outputs.logits.squeeze(-1).float().detach().numpy()
+score = logits.item()
+result = {
+    "text": text,
+    "score": score,
+    "int_score": int(round(max(0, min(score, 5)))),
+}
 
-More information needed
+print(result)
+# {'text': 'This is a test sentence.', 'score': 0.07964489609003067, 'int_score': 0}
+```
 
 ## Intended uses & limitations
 
-More information needed
+While the classifier performs well in distinguishing high-quality code in its target language (Ruby in this case), there are some limitations:
 
-## Training and evaluation data
+- Scope: The model's performance might change for other datasets, in particular for out of distribution samples. The classifier's context is 1024 tokens, which might not be sufficient to assess the quality of some long code files.
+- Bias: The model's performance is dependent on the quality and representativeness of the training data and the LLM used for the annotation. Biases in both can affect the classifier's judgments. It might overfit to thoroughly commented code.
+- Context: The classifier evaluates individual code files without considering broader context, which might impact its effectiveness in certain scenarios.
 
-More information needed
+The training and inference code is available on GitHub 
+https://github.com/huggingface/cosmopedia/tree/main/classification
 
 ## Training procedure
 
+The classifier was trained on 500,000 pairs of code files and their scores from 0 to 5, generated by Llama3.1. The samples were annotated based on their educational quality with 1 being not educational and 5 being highly educational and relevant for teaching programming. You can find the prompt used for building the annotations in the appendix of [SmolLM2 paper](https://arxiv.org/pdf/2502.02737).
+
+We added a classification head with a single regression output to StarEncoder and trained the model for 20 epochs with a learning rate of 3e-4. During training, the embedding and encoder layers were frozen to focus on the classification head. 
+
+It achieves the following results on the evaluation set:
+- Loss: 0.3282
+- Precision: 0.4623
+- Recall: 0.3260
+- F1 Macro: 0.3536
+- Accuracy: 0.6657
+- F1 Binary Minimum3: 0.6101
+
+While the macro F1 scores across the 1-5 rating scale are relatively low due to the model's difficulty in distinguishing between higher-rated samples, the classifier performs well for our primary filtering task. When converting to binary classification, using a threshold of 2 achieves the F1 scores ranges between 0.8 and 0.9 for most Stack-Edu  classifiers, whereas a threshold of 3 yields F1 scores between 0.55 and 0.70. With the Highest being Python, C, Rust and the lowest being TypeScript and PHP.
+
+<div style="display: flex; justify-content: center; gap: 20px;">
+  <img src="https://cdn-uploads.huggingface.co/production/uploads/61c141342aac764ce1654e43/z0lCnH2nqYMu09JsxPvFb.png" width="600">
+  <img src="https://cdn-uploads.huggingface.co/production/uploads/61c141342aac764ce1654e43/W6BQO6Rkafpj-96fWdCOv.png" width="600">
+</div>
+We validated these classifiers by filtering Stack v2 data and testing on an intermediate SmolLM2 checkpoint. Filtering with a threshold of 3 improved performance across most languages while maintaining adequate data volume, though Java showed better results with a threshold of 2.
+
 ### Training hyperparameters
 
 The following hyperparameters were used during training:
@@ -55,29 +88,19 @@ The following hyperparameters were used during training:
 - lr_scheduler_warmup_steps: 200
 - num_epochs: 20
 
-### Training results
-
-| Training Loss | Epoch   | Step  | Validation Loss | Precision | Recall | F1 Macro | Accuracy | F1 Binary Minimum3 |
-|:-------------:|:-------:|:-----:|:---------------:|:---------:|:------:|:--------:|:--------:|:------------------:|
-| No log        | 0       | 0     | 4.7647          | 0.0010    | 0.1667 | 0.0020   | 0.0062   | 0                  |
-| 0.358         | 1.4368  | 1000  | 0.3549          | 0.4093    | 0.2953 | 0.3072   | 0.6513   | 0.5882             |
-| 0.3559        | 2.8736  | 2000  | 0.3425          | 0.4649    | 0.3087 | 0.3294   | 0.6571   | 0.6143             |
-| 0.3534        | 4.3103  | 3000  | 0.3391          | 0.4318    | 0.3144 | 0.3349   | 0.6586   | 0.6149             |
-| 0.3539        | 5.7471  | 4000  | 0.3394          | 0.4219    | 0.3244 | 0.3446   | 0.6579   | 0.6298             |
-| 0.3585        | 7.1839  | 5000  | 0.3359          | 0.4756    | 0.3106 | 0.3350   | 0.6622   | 0.6069             |
-| 0.3476        | 8.6207  | 6000  | 0.3339          | 0.4551    | 0.3178 | 0.3415   | 0.6638   | 0.6082             |
-| 0.3496        | 10.0575 | 7000  | 0.3307          | 0.4512    | 0.3263 | 0.3505   | 0.6656   | 0.6204             |
-| 0.3362        | 11.4943 | 8000  | 0.3307          | 0.4657    | 0.3228 | 0.3485   | 0.6640   | 0.6178             |
-| 0.3442        | 12.9310 | 9000  | 0.3307          | 0.4771    | 0.3248 | 0.3517   | 0.6677   | 0.6095             |
-| 0.344         | 14.3678 | 10000 | 0.3287          | 0.4774    | 0.3222 | 0.3496   | 0.6660   | 0.6147             |
-| 0.3332        | 15.8046 | 11000 | 0.3281          | 0.4678    | 0.3240 | 0.3504   | 0.6658   | 0.6168             |
-| 0.3359        | 17.2414 | 12000 | 0.3300          | 0.4658    | 0.3203 | 0.3471   | 0.6643   | 0.6100             |
-| 0.3306        | 18.6782 | 13000 | 0.3282          | 0.4623    | 0.3260 | 0.3536   | 0.6657   | 0.6101             |
-
-
-### Framework versions
-
-- Transformers 4.43.4
-- Pytorch 2.4.0+cu121
-- Datasets 2.21.0
-- Tokenizers 0.19.1
+## License
+
+[Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0)
+
+## Citation
+```bash
+@misc{allal2025smollm2smolgoesbig,
+      title={SmolLM2: When Smol Goes Big -- Data-Centric Training of a Small Language Model}, 
+      author={Loubna Ben Allal and Anton Lozhkov and Elie Bakouch and Gabriel Martín Blázquez and Guilherme Penedo and Lewis Tunstall and Andrés Marafioti and Hynek Kydlíček and Agustín Piqueres Lajarín and Vaibhav Srivastav and Joshua Lochner and Caleb Fahlgren and Xuan-Son Nguyen and Clémentine Fourrier and Ben Burtenshaw and Hugo Larcher and Haojun Zhao and Cyril Zakka and Mathieu Morlon and Colin Raffel and Leandro von Werra and Thomas Wolf},
+      year={2025},
+      eprint={2502.02737},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2502.02737}, 
+}
+```