Update README.md

README.md

@@ -8,17 +8,22 @@ license: apache-2.0
 
 # Model Overview
 This is a text classification model to classify documents into one of 26 domain classes:
-
+```
 'Adult', 'Arts_and_Entertainment', 'Autos_and_Vehicles', 'Beauty_and_Fitness', 'Books_and_Literature', 'Business_and_Industrial', 'Computers_and_Electronics', 'Finance', 'Food_and_Drink', 'Games', 'Health', 'Hobbies_and_Leisure', 'Home_and_Garden', 'Internet_and_Telecom', 'Jobs_and_Education', 'Law_and_Government', 'News', 'Online_Communities', 'People_and_Society', 'Pets_and_Animals', 'Real_Estate', 'Science', 'Sensitive_Subjects', 'Shopping', 'Sports', 'Travel_and_Transportation'
+```
+
 # Model Architecture
-The model architecture is Deberta V3 Base
-Context length is 512 tokens
-# Training (details)
+- The model architecture is Deberta V3 Base
+- Context length is 512 tokens
+
+# Training Details
 ## Training data:
 - 1 million Common Crawl samples, labeled using Google Cloud’s Natural Language API: https://cloud.google.com/natural-language/docs/classifying-text
 - 500k Wikepedia articles, curated using Wikipedia-API: https://pypi.org/project/Wikipedia-API/
+
 ## Training steps:
 Model was trained in multiple rounds using Wikipedia and Common Crawl data, labeled by a combination of pseudo labels and Google Cloud API.
+
 # How To Use This Model
 ## Input
 The model takes one or several paragraphs of text as input.
@@ -38,15 +43,14 @@ Example output:
 Food_and_Drink
 ```
 
-# How to use in NeMo Curator
+# How to Use in NVIDIA NeMo Curator
 
-The inference code is available on [NeMo Curator's GitHub repository](https://github.com/NVIDIA/NeMo-Curator). Download the [model.pth](https://huggingface.co/nvidia/domain-classifier/blob/main/model.pth) file and check out this [example notebook](https://github.com/NVIDIA/NeMo-Curator/blob/main/tutorials/distributed_data_classification/distributed_data_classification.ipynb) to get started.
+The inference code is available on [NeMo Curator's GitHub repository](https://github.com/NVIDIA/NeMo-Curator). Check out this [example notebook](https://github.com/NVIDIA/NeMo-Curator/blob/main/tutorials/distributed_data_classification/distributed_data_classification.ipynb) to get started.
 
-# How to use in transformers
-To use the Domain classifier, use the following code:
+# How to Use in Transformers
+To use the domain classifier, use the following code:
 
 ```python
-
 import torch
 from torch import nn
 from transformers import AutoModel, AutoTokenizer, AutoConfig
@@ -55,9 +59,9 @@ from huggingface_hub import PyTorchModelHubMixin
 class CustomModel(nn.Module, PyTorchModelHubMixin):
     def __init__(self, config):
         super(CustomModel, self).__init__()
-        self.model = AutoModel.from_pretrained(config['base_model'])
-        self.dropout = nn.Dropout(config['fc_dropout'])
-        self.fc = nn.Linear(self.model.config.hidden_size, len(config['id2label']))
+        self.model = AutoModel.from_pretrained(config["base_model"])
+        self.dropout = nn.Dropout(config["fc_dropout"])
+        self.fc = nn.Linear(self.model.config.hidden_size, len(config["id2label"]))
 
     def forward(self, input_ids, attention_mask):
         features = self.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
@@ -73,7 +77,7 @@ model = CustomModel.from_pretrained("nvidia/domain-classifier")
 # Prepare and process inputs
 text_samples = ["Sports is a popular domain", "Politics is a popular domain"]
 inputs = tokenizer(text_samples, return_tensors="pt", padding="longest", truncation=True)
-outputs = model(inputs['input_ids'], inputs['attention_mask'])
+outputs = model(inputs["input_ids"], inputs["attention_mask"])
 
 # Predict and display results
 predicted_classes = torch.argmax(outputs, dim=1)
@@ -122,6 +126,7 @@ PR-AUC score for each domain:
 # References
 - https://arxiv.org/abs/2111.09543
 - https://github.com/microsoft/DeBERTa
+
 # License
 License to use this model is covered by the Apache 2.0. By downloading the public and release version of the model, you accept the terms and conditions of the Apache License 2.0.
-This repository contains the code for the domain classifier model.
+This repository contains the code for the domain classifier model.