Update README.md

README.md

@@ -1,26 +1,34 @@
-# Suicidal-BERT
-This text classification model predicts whether a sequence of words are suicidal (1) or non-suicidal (0).
-
-## Data
-The model was trained on the [Suicide and Depression Dataset](https://www.kaggle.com/nikhileswarkomati/suicide-watch) obtained from Kaggle. The dataset was scraped from Reddit and consists of 232,074 rows equally distributed between 2 classes - suicide and non-suicide.
-
-## Parameters
-The model fine-tuning was conducted on 1 epoch, with batch size of 6, and learning rate of 0.00001. Due to limited computing resources and time, we were unable to scale up the number of epochs and batch size.
-
-## Performance
-The model has achieved the following results after fine-tuning on the aforementioned dataset:
-- Accuracy: 0.9757
-- Recall: 0.9669
-- Precision: 0.9701
-- F1 Score: 0.9685
-
-## How to Use
-Load the model via the transformers library:
-```
-from transformers import AutoTokenizer, AutoModel
-tokenizer = AutoTokenizer.from_pretrained("gooohjy/suicidal-bert")
-model = AutoModel.from_pretrained("gooohjy/suicidal-bert")
-```
-
-## Resources
+---
+license: mit
+language:
+- en
+base_model:
+- google-bert/bert-base-uncased
+pipeline_tag: text-classification
+---
+# Suicidal-BERT
+This text classification model predicts whether a sequence of words are suicidal (1) or non-suicidal (0).
+
+## Data
+The model was trained on the [Suicide and Depression Dataset](https://www.kaggle.com/nikhileswarkomati/suicide-watch) obtained from Kaggle. The dataset was scraped from Reddit and consists of 232,074 rows equally distributed between 2 classes - suicide and non-suicide.
+
+## Parameters
+The model fine-tuning was conducted on 1 epoch, with batch size of 6, and learning rate of 0.00001. Due to limited computing resources and time, we were unable to scale up the number of epochs and batch size.
+
+## Performance
+The model has achieved the following results after fine-tuning on the aforementioned dataset:
+- Accuracy: 0.9757
+- Recall: 0.9669
+- Precision: 0.9701
+- F1 Score: 0.9685
+
+## How to Use
+Load the model via the transformers library:
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained("gooohjy/suicidal-bert")
+model = AutoModel.from_pretrained("gooohjy/suicidal-bert")
+```
+
+## Resources
 For more resources, including the source code, please refer to the GitHub repository [gohjiayi/suicidal-text-detection](https://github.com/gohjiayi/suicidal-text-detection/).