Update README.md

README.md

@@ -5,7 +5,7 @@ metrics:
 ---
 # BERT Text Classification Model
 
-This is a simple demo model for text classification using BERT.
+This is a simple model for text classification using BERT.
 
 ## Usage
 
@@ -16,4 +16,22 @@ text = "This is a positive review."
 predicted_class = classify_text(text)
 print("Predicted class:", predicted_class)
 
-Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2018). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv preprint arXiv:1810.04805.
+from transformers import BertTokenizer, BertForSequenceClassification
+
+# Load pre-trained BERT tokenizer and model
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+
+# Define a function to classify text
+def classify_text(text):
+    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True)
+    outputs = model(**inputs)
+    logits = outputs.logits
+    probabilities = logits.softmax(dim=1)
+    predicted_class = probabilities.argmax(dim=1).item()
+    return predicted_class
+
+# Example usage
+text = "This is a positive review."
+predicted_class = classify_text(text)
+print("Predicted class:", predicted_class)