Update app.py

app.py

@@ -1,6 +1,21 @@
-# Use a pipeline as a high-level helper
-from transformers import pipeline
+import torch
+from transformers import AutoTokenizer, BertForSequenceClassification
 
-pipe = pipeline("text-classification", model="priyabrat/bert_categorisation")
+tokenizer = AutoTokenizer.from_pretrained("textattack/bert-base-uncased-yelp-polarity")
+model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-yelp-polarity")
 
-)
+inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+
+with torch.no_grad():
+    logits = model(**inputs).logits
+
+predicted_class_id = logits.argmax().item()
+model.config.id2label[predicted_class_id]
+
+# To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
+num_labels = len(model.config.id2label)
+model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-yelp-polarity", num_labels=num_labels)
+
+labels = torch.tensor([1])
+loss = model(**inputs, labels=labels).loss
+round(loss.item(), 2)