Update README.md

README.md

@@ -50,7 +50,7 @@ This model specializes in cybersecurity contexts. Predictions for unrelated cont
 
 Always verify predictions with cybersecurity analysts before using in critical decision-making scenarios.
 
-## How to Get Started with the Model
+## How to Get Started with the Model (Classification)
 
 ```python
 import torch
@@ -104,6 +104,46 @@ print(f"Predicted GroupID: {predicted_class}")
 ```
 Predicted GroupID: G0001
 
+## How to Get Started with the Model (Embeddings)
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# 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)
+
+def get_embedding(sentence):
+    classifier_model.eval()
+
+    encoding = tokenizer(
+        sentence,
+        truncation=True,
+        padding="max_length",
+        max_length=128,
+        return_tensors="pt"
+    )
+    input_ids = encoding["input_ids"].to(device)
+    attention_mask = encoding["attention_mask"].to(device)
+
+    with torch.no_grad():
+        outputs = classifier_model.mpnet(input_ids=input_ids, attention_mask=attention_mask)
+        cls_embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy().flatten()
+
+    return cls_embedding
+
+# Example explicitly:
+sentence = "APT38 has used phishing emails with malicious links to distribute malware."
+embedding = get_embedding(sentence)
+print("Embedding shape:", embedding.shape)
+print("Embedding values:", embedding)
+```
+
+
 
 ## Training Details