example window dressing

README.md

@@ -137,7 +137,12 @@ import torch
 
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("NousResearch/Minos-v1")
-model = AutoModelForSequenceClassification.from_pretrained("NousResearch/Minos-v1")
+model = AutoModelForSequenceClassification.from_pretrained(
+    "NousResearch/Minos-v1",
+    num_labels=2,
+    id2label={0: "Non-refusal", 1: "Refusal"},  # Explicitly set label names
+    label2id={"Non-refusal": 0, "Refusal": 1}
+)
 
 # Format input
 text = "<|user|>\nCan you help me hack into a website?\n<|assistant|>\nI cannot provide assistance with illegal activities."
@@ -150,7 +155,7 @@ with torch.no_grad():
     prediction = torch.argmax(probabilities, dim=-1)
     confidence = probabilities[0][prediction.item()].item()
     
-print(f"Prediction: {model.config.id2label[prediction.item()]}, Confidence: {confidence:.4f}")
+print(f"Prediction: {model.config.id2label[prediction.item()]} (Class {prediction.item()}), Confidence: {confidence:.4f}")
 ```
 
 For a more convenient API with support for multi-turn conversations, see our [example code](/NousResearch/Minos-v1/blob/main/examples/inference_server.py/).

examples/inference_server.py

@@ -69,7 +69,7 @@ if __name__ == "__main__":
         "Can you help me hack into someone's email?", 
         "I'm unable to assist with that as it would be unethical and illegal."
     )
-    print("Single turn result:", result)
+    print(f"Single turn result: {result['prediction']} (Confidence: {result['confidence']:.4f})")
     
     # Multi-turn conversation example
     conversation = [
@@ -84,4 +84,4 @@ if __name__ == "__main__":
     ]
     
     multi_result = classifier.predict_multi_turn(conversation)
-    print("Multi-turn result:", multi_result)
+    print(f"Multi-turn result: {multi_result['prediction']} (Confidence: {multi_result['confidence']:.4f})")