Update README.md

README.md

@@ -62,11 +62,6 @@ from huggingface_hub import hf_hub_download
 import json
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# Load explicitly your fine-tuned MPNet model
-classifier_model = AutoModelForSequenceClassification.from_pretrained("selfconstruct3d/AttackGroup-MPNET").to(device)
-
-# Load explicitly your tokenizer
-tokenizer = AutoTokenizer.from_pretrained("selfconstruct3d/AttackGroup-MPNET")
 
 
 label_to_groupid_file = hf_hub_download(
@@ -77,6 +72,12 @@ label_to_groupid_file = hf_hub_download(
 with open(label_to_groupid_file, "r") as f:
     label_to_groupid = json.load(f)
 
+# Load explicitly your fine-tuned MPNet model
+classifier_model = AutoModelForSequenceClassification.from_pretrained("selfconstruct3d/AttackGroup-MPNET", num_labels=len(label_to_groupid)).to(device)
+
+# Load explicitly your tokenizer
+tokenizer = AutoTokenizer.from_pretrained("selfconstruct3d/AttackGroup-MPNET")
+
 def predict_group(sentence):
     classifier_model.eval()
     encoding = tokenizer(
@@ -109,13 +110,25 @@ Predicted GroupID: G0001
 ```python
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from huggingface_hub import hf_hub_download
+import json
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+
+label_to_groupid_file = hf_hub_download(
+    repo_id="selfconstruct3d/AttackGroup-MPNET",
+    filename="label_to_groupid.json"
+)
+
+with open(label_to_groupid_file, "r") as f:
+    label_to_groupid = json.load(f)
+
+
 # Load your fine-tuned classification model
 model_name = "selfconstruct3d/AttackGroup-MPNET"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-classifier_model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
+classifier_model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(label_to_groupid)).to(device)
 
 def get_embedding(sentence):
     classifier_model.eval()