Update README.md

updated README.md

README.md

@@ -67,11 +67,38 @@ import torch
 model = AutoModelForSequenceClassification.from_pretrained("mozilla/content-multilabel-iab-classifier")
 tokenizer = AutoTokenizer.from_pretrained("mozilla/content-multilabel-iab-classifier")
 
+label_list = [
+    'inconclusive',
+    'animals',
+    'arts',
+    'autos',
+    'business',
+    'career',
+    'education',
+    'fashion',
+    'finance',
+    'food',
+    'government',
+    'health',
+    'hobbies',
+    'home',
+    'news',
+    'realestate',
+    'society',
+    'sports',
+    'tech',
+    'travel'
+]
+label2id = {label: idx for idx, label in enumerate(label_list)}
+id2label = {idx: label for label, idx in label2id.items()}
+
 text = "Discover the latest trends in AI and wearable technology."
-inputs = tokenizer(text, return_tensors="pt", truncation=True)
-outputs = model(**inputs)
-probs = torch.sigmoid(outputs.logits)
-predicted_labels = (probs > 0.5).int()
+
+with torch.no_grad():
+  inputs = tokenizer(text, return_tensors="pt", truncation=True)
+  outputs = model(**inputs)
+  probs = torch.sigmoid(outputs.logits).squeeze().cpu().numpy()
+  predicted_labels = [(id2label[i], round(p, 3)) for i, p in enumerate(probs) if p >= 0.5]
 ```
 
 📖 Citation