Add PyTorchModelHubMixin

README.md

@@ -81,9 +81,8 @@ To use this AEGIS classifiers, you must get access to Llama Guard on Hugging Fac
 ```python
 import torch
 import torch.nn.functional as F
-from huggingface_hub import hf_hub_download
+from huggingface_hub import PyTorchModelHubMixin
 from peft import PeftModel
-from safetensors.torch import load_file
 from torch.nn import Dropout, Linear
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -103,8 +102,8 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 tokenizer.pad_token = tokenizer.unk_token
 
-class InstructionDataGuardNet(torch.nn.Module):
-    def __init__(self, input_dim, dropout=0.7):
+class InstructionDataGuardNet(torch.nn.Module, PyTorchModelHubMixin):
+    def __init__(self, input_dim=4096, dropout=0.7):
         super().__init__()
         self.input_dim = input_dim
         self.dropout = Dropout(dropout)
@@ -129,13 +128,8 @@ class InstructionDataGuardNet(torch.nn.Module):
         return x
 
 # Load Instruction-Data-Guard classifier
-instruction_data_guard = InstructionDataGuardNet(4096).to(device)
-weights_path = hf_hub_download(
-    repo_id="nvidia/instruction-data-guard",
-    filename="model.safetensors",
-)
-state_dict = load_file(weights_path)
-instruction_data_guard.load_state_dict(state_dict)
+instruction_data_guard = InstructionDataGuardNet.from_pretrained("nvidia/instruction-data-guard")
+instruction_data_guard = instruction_data_guard.to(device)
 instruction_data_guard = instruction_data_guard.eval()
 
 # Function to compute results