Update README.md

README.md

@@ -8,7 +8,6 @@ datasets:
 language:
 - en
 pipeline_tag: text-classification
-library_name: transformers
 ---
 # Sample AG News Classifier
 
@@ -22,28 +21,79 @@ The model has been trained to classify news headlines into 4 categories: World,
 Note: Currently limited to above 4 categories only
 
 # How to use
-You can use the model using pipeline or AutoModelForMultimodalLM options offered by HuggingFace
+The model has been developed using PyTorch and since it is a custom neural network architecture, it can be used by performing the below steps and not using the Transformers library,
 
+1. Create your model class
+ 
 ```python
-# Use a pipeline as a high-level helper
-from transformers import pipeline
+import torch
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
 
-pipe = pipeline("text-classification", model="pulkitchowdry/sample-agnews-classifer")
-text = "US-Iran sign a peace deal"
-pipe(text)
 
-# Load model directly
-from transformers import AutoTokenizer, AutoModel
+class TextClassifier(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, vocab_size=30522, embed_dim=128, num_classes=4):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.fc1 = nn.Linear(embed_dim, 128)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(128, num_classes)
 
-tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
-model = AutoModel.from_pretrained("pulkitchowdry/sample-agnews-classifer", dtype="auto")
+    def forward(self, input_ids, attention_mask=None):
+        x = self.embedding(input_ids)
 
-text = "US-Iran sign a peace deal"
-inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=128)
-prediction = model(inputs)
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).float()
+            x = x * mask
+            x = x.sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
+        else:
+            x = x.mean(dim=1)
+
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        return x
+
+model = TextClassifier.from_pretrained("pulkitchowdry/sample-agnews-classifer")
+model.eval()
+
+```
+
+2. Setup Tokenizer
+
+```python
+
+from transformers import AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+text      = "Messi scores a hatrick in the world cup"
+
+inputs  = tokenizer(
+                text, 
+                return_tensors="pt", 
+                padding="max_length", 
+                truncation=True, 
+                max_length=128
+              )
+
+```
+
+3. Run the prediction
+
+```python
+
+with torch.no_grad():
+  logits  = model(
+      inputs["input_ids"],
+      attention_mask=inputs["attention_mask"]
+  )
+  prediction  = logits.argmax(dim=1).item()
+
+print(f"{prediction}")
 categories = ["World", "Sports", "Business", "Science/Technology"]
 
-print(categories[prediction])
+print("Predicted class: ", categories[prediction])
+
 ```
 
 # Limitations and bias