Update README.md

README.md

@@ -47,24 +47,32 @@ pip install torch transformers safetensors
 Load and run the model using PyTorch and transformers:
 
 ```python
-import torch
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoModelForTokenClassification, AutoConfig, BertTokenizerFast
 from safetensors.torch import load_file
 
+# Load the config
+config = AutoConfig.from_pretrained("folder_to_model")
+
+# Initialize the model with the config
+model = AutoModelForTokenClassification.from_config(config)
+
+# Load the safetensors weights
+state_dict = load_file("folder_to_tensors")
+
+# Load the state dict into the model
+model.load_state_dict(state_dict)
+
 # Load the tokenizer
 tokenizer = BertTokenizerFast.from_pretrained("google-bert/bert-base-multilingual-cased")
 
-# Load the model
-model = AutoModel.from_pretrained('model-path/miniagent.pt', device_map='auto')
-# Alternatively, for the precision-focused model
-# model = AutoModel.from_pretrained('model-path/miniagent_precision', device_map='auto')
-
-# Example input
-text = "Your sensitive information string"
+# Load the label mapper if needed
+with open("pii_model/label_mapper.json", 'r') as f:
+    label_mapper_data = json.load(f)
 
-# Tokenize and run the model
-inputs = tokenizer(text, return_tensors="pt")
-outputs = model(**inputs)
+label_mapper = LabelMapper()
+label_mapper.label_to_id = label_mapper_data['label_to_id']
+label_mapper.id_to_label = {int(k): v for k, v in label_mapper_data['id_to_label'].items()}
+label_mapper.num_labels = label_mapper_data['num_labels']
 
 # Process outputs for analysis...
 ```