multilingual fine-edu classifier

README.md

@@ -1,9 +1,25 @@
+language:
+- multilingual
+license: apache-2.0
+---
 
+## Running Model:
+To run inference you must install
+```
+pip install transformers[torch]
+pip install datasets
+pip install pandas
+pip install tqdm
+```
+
+After installing those libraries you can sun the following code:
 
 ```python
+import pandas as pd
 import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
-from datasets import Dataset
+from tqdm import tqdm
+
 
 device = "cuda"
 path = "Unbabel/mfineweb-edu-classifier"
@@ -15,19 +31,104 @@ model = AutoModelForSequenceClassification.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True)
 
-def tokenize(examples):
-    return tokenizer(examples["text"], truncation=True, max_length=512)
+def get_model_outputs(texts):
+    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512).to(model.device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        score = outputs.logits
+        prob = torch.nn.functional.sigmoid(outputs.binary_logits)
+    return score.cpu(), prob.cpu()
+
+def batchify_texts(texts, batch_size):
+    for i in range(0, len(texts), batch_size):
+        yield texts[i:i + batch_size]
+
+# TODO: replace the next line with the texts you want to classify
+texts = LIST_WITH_TEXTS_TO_CLASSIFY
+batch_size = 64  # Adjust based on your available memory and model capacity
+num_batches = (len(texts) + batch_size - 1) // batch_size
+
+all_scores = []
+all_probs = []
+with tqdm(total=num_batches, dynamic_ncols=True) as pbar:
+    for batch_num, batch in enumerate(batchify_texts(texts, batch_size), 1):
+        score, probs = get_model_outputs(batch)
+        all_scores.append(score)
+        all_probs.append(probs)
+        pbar.set_description(f"Processing Batch {batch_num}/{num_batches}")
+        pbar.update(1)
+
+# SCORES is the output of the regression head and should reflect the
+# educational score of the text!
+scores = torch.cat(all_scores, dim=0).squeeze()
+
+## BINARY_PRED is the output of the classification head that tells
+# if a text has an acceptable educational score or not.
+# NOTE: Converting the scores into binary predictions is also possible
+all_probs = torch.cat(all_probs, dim=0).squeeze()
+binary_pred = (all_probs >= 0.5).numpy().astype(int)
+```
+
+## English Results:
+
+When testing the model on an english partition with 37537 samples the results are comparable to the original FineEdu-classifier.
+
+Regression head results:
+```
+              precision    recall  f1-score   support
+
+           0       0.80      0.53      0.64      5130
+           1       0.80      0.88      0.83     21602
+           2       0.63      0.58      0.61      7849
+           3       0.54      0.62      0.58      2310
+           4       0.62      0.48      0.54       645
+           5       0.00      0.00      0.00         1
+
+    accuracy                           0.74     37537
+   macro avg       0.56      0.51      0.53     37537
+weighted avg       0.74      0.74      0.74     37537
+```
+
+Binary head results:
+```
+              precision    recall  f1-score   support
+
+           0       0.98      0.97      0.98     34581
+           1       0.71      0.74      0.73      2956
+
+    accuracy                           0.96     37537
+   macro avg       0.85      0.86      0.85     37537
+weighted avg       0.96      0.96      0.96     37537
+```
+
+## Multilingual Results:
+
+If we evaluate on the same texts translated into 15 different languages are almost identical!
+
+Regression head results:
+```
+              precision    recall  f1-score   support
+
+           0       0.80      0.50      0.61      5130
+           1       0.79      0.87      0.83     21602
+           2       0.61      0.58      0.59      7849
+           3       0.52      0.61      0.56      2310
+           4       0.61      0.38      0.47       645
+           5       0.00      0.00      0.00         1
+
+    accuracy                           0.73     37537
+   macro avg       0.55      0.49      0.51     37537
+weighted avg       0.73      0.73      0.73     37537
+```
 
-texts = [
-    "This is a text",
-    "this is another text to classify"
-]
+Binary head results:
+```
+              precision    recall  f1-score   support
 
-model_inputs = [
-    tokenizer(text, truncation=True, max_length=512) for text in texts
-]
+           0       0.98      0.97      0.97     34581
+           1       0.70      0.71      0.71      2956
 
-with torch.no_grad():
-    for model_input in model_inputs:
-        output = model(input_ids)
-```
+    accuracy                           0.95     37537
+   macro avg       0.84      0.84      0.84     37537
+weighted avg       0.95      0.95      0.95     37537
+```

train.py

@@ -0,0 +1,76 @@
+import numpy as np
+import pandas as pd
+from classifier import DebertaV2ForSequenceClassification
+from datasets import Dataset
+from scipy.stats import pearsonr
+from sklearn.metrics import accuracy_score, precision_score, recall_score
+from transformers import (AutoTokenizer, DataCollatorWithPadding, Trainer,
+                          TrainingArguments)
+
+tokenizer = AutoTokenizer.from_pretrained("microsoft/mdeberta-v3-base")
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+def compute_metrics(eval_pred):
+    predictions, labels = eval_pred
+    scores, binary_logits = predictions
+    scores = scores.squeeze()
+    probs = sigmoid(binary_logits.squeeze())
+    predicted_labels = (probs >= 0.5).astype(int)
+    binary_labels = (labels >= 3).astype(int)
+    return {
+        'pearson': pearsonr(scores, labels)[0],
+        'accuracy': accuracy_score(binary_labels, predicted_labels),
+        'precision': precision_score(binary_labels, predicted_labels),
+        'recall': recall_score(binary_labels, predicted_labels),
+    }
+
+def tokenize_function(examples):
+    return tokenizer(examples["text"], truncation=True, max_length=512)
+
+def train_classifier():
+    train_csv = pd.read_csv(PATH_TO_TRAINSET)
+    train_dataset = Dataset.from_pandas(train_csv)
+    
+    test_csv = pd.read_csv(PATH_TO_TESTSET).sample(n=10_000, random_state=42)
+    test_dataset = Dataset.from_pandas(test_csv)
+    
+    train_dataset = train_dataset.map(tokenize_function, batched=True)
+    test_dataset = test_dataset.map(tokenize_function, batched=True)
+    train_dataset = train_dataset.with_format("torch")
+    test_dataset = test_dataset.with_format("torch")
+    
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
+    training_args = TrainingArguments(
+        output_dir="./results",
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        learning_rate=2e-5,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        num_train_epochs=3,
+        weight_decay=0.01,
+        logging_dir="./logs",
+        logging_steps=10,
+    )
+    model = DebertaV2ForSequenceClassification.from_pretrained("microsoft/mdeberta-v3-base")
+    print ("Freezing model embeddings!")
+    model.freeze_embeddings()
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=test_dataset,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics
+    )
+    trainer.train()
+    # Evaluate the model
+    trainer.evaluate()
+    #trainer.push_to_hub(private=True, model_name="mFine-Edu-classifier")
+
+if __name__ == "__main__":
+    train_classifier()