Update README.md

README.md

@@ -11,38 +11,107 @@ base_model:
 - google-bert/bert-base-uncased
 library_name: transformers
 ---
-# BERT Fine-Tuned for Frugal AI Challenge - Text Task
+# Model Card: BERT Fine-Tuned for Frugal AI Challenge - Text Task
 
 ## Model Overview
-This model is a fine-tuned version of `bert-base-uncased` for multi-class text classification. It was trained on the Frugal AI Challenge Text Task dataset.
+This model is a fine-tuned version of the `bert-base-uncased` transformer model, specifically tailored for multi-class text classification. It was developed as part of the Frugal AI Challenge to classify text into eight distinct categories. The model incorporates a custom classification head and leverages class weighting to address dataset imbalance.
+
+## Dataset
+- **Source**: Frugal AI Challenge Text Task Dataset
+- **Classes**: 8 unique labels representing various categories of text
+- **Preprocessing**: Tokenization using `BertTokenizer` with padding and truncation to a maximum sequence length of 128.
+
+## Model Architecture
+- **Base Model**: `bert-base-uncased`
+- **Classification Head**: Custom head with weighted cross-entropy loss to handle class imbalance.
+- **Number of Labels**: 8
 
 ## Training Details
-- **Dataset**: Frugal AI Challenge Text Task
-- **Architecture**: BERT with a custom classification head
 - **Optimizer**: AdamW
 - **Learning Rate**: 2e-5
+- **Batch Size**: 16 (for both training and evaluation)
 - **Epochs**: 3
-- **Batch Size**: 16
+- **Weight Decay**: 0.01
+- **Evaluation Strategy**: Performed at the end of each epoch
+- **Hardware**: Trained on GPUs for efficient computation
 
 ## Performance Metrics (Validation Set)
-| Metric       | Value       |
-|--------------|-------------|
-| Accuracy     | 68          |
-| Precision    | 69          |
-| Recall       | 68          |
-| F1-Score     | 68          |
+The model achieved the following performance metrics on the validation set:
+
+| Class                          | Precision | Recall | F1-Score | Support |
+|--------------------------------|-----------|--------|----------|---------|
+| 5_science_unreliable           | 0.65      | 0.71   | 0.68     | 160     |
+| 1_not_happening                | 0.69      | 0.80   | 0.74     | 148     |
+| 4_solutions_harmful_unnecessary| 0.62      | 0.72   | 0.66     | 155     |
+| 0_not_relevant                 | 0.82      | 0.60   | 0.70     | 324     |
+| 6_proponents_biased            | 0.62      | 0.64   | 0.63     | 157     |
+| 7_fossil_fuels_needed          | 0.59      | 0.68   | 0.63     | 57      |
+| 2_not_human                    | 0.68      | 0.70   | 0.69     | 141     |
+| 3_not_bad                      | 0.59      | 0.62   | 0.61     | 77      |
+
+- **Overall Accuracy**: 68%
+- **Macro Average**: Precision: 0.66, Recall: 0.68, F1-Score: 0.67
+- **Weighted Average**: Precision: 0.69, Recall: 0.68, F1-Score: 0.68
 
 ## Training Evolution
 ### Training and Validation Loss
+The training and validation loss evolution over epochs is shown below:
+
 ![Training Loss](./training_loss_plot.png)
 
 ### Validation Accuracy
+The validation accuracy evolution over epochs is shown below:
+
 ![Validation Accuracy](./validation_accuracy_plot.png)
 
-## Classification Report (Validation Set)
+## Confusion Matrix
+The confusion matrix below illustrates the model's performance on the validation set, highlighting areas of strength and potential misclassifications:
 
+![Confusion Matrix](./confusion_matrix.png)
 
-## Confusion Matrix
-The confusion matrix below shows the model's performance on the validation set.
+## Key Features
+- **Class Weighting**: Addressed dataset imbalance by incorporating class weights during training.
+- **Custom Loss Function**: Used weighted cross-entropy loss for better handling of underrepresented classes.
+- **Evaluation Metrics**: Accuracy, precision, recall, and F1-score were computed to provide a comprehensive understanding of the model's performance.
+
+## Usage
+This model can be used for multi-class text classification tasks where the input text needs to be categorized into one of the eight predefined classes. It is particularly suited for datasets with class imbalance, thanks to its weighted loss function.
+
+### Example Usage
+```python
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+# Load the fine-tuned model and tokenizer
+model = AutoModelForSequenceClassification.from_pretrained("ParisNeo/bert-frugal-ai-text-classification")
+tokenizer = AutoTokenizer.from_pretrained("ParisNeo/bert-frugal-ai-text-classification")
+
+# Tokenize input text
+text = "Your input text here"
+inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=128)
+
+# Perform inference
+outputs = model(**inputs)
+predicted_class = outputs.logits.argmax(-1).item()
+
+print(f"Predicted Class: {predicted_class}")
+```
+
+## Limitations
+- **Dataset-Specific**: The model's performance is optimized for the Frugal AI Challenge dataset and may require further fine-tuning for other datasets.
+- **Class Imbalance**: While class weighting mitigates imbalance, some underrepresented classes may still have lower performance.
+- **Sequence Length**: Input text is truncated to a maximum length of 128 tokens, which may result in loss of information for longer texts.
+
+## Citation
+If you use this model in your research or application, please cite it as:
+```
+@model{ParisNeo_bert_frugal_ai_text_classification,
+  author = {ParisNeo},
+  title = {BERT Fine-Tuned for Frugal AI Challenge - Text Task},
+  year = {2025},
+  publisher = {Hugging Face},
+  url = {https://huggingface.co/ParisNeo/bert-frugal-ai-text-classification}
+}
+```
 
-![Confusion Matrix](./confusion_matrix.png)
+## Acknowledgments
+Special thanks to the Frugal AI Challenge organizers for providing the dataset and fostering innovation in AI research.