Update README.md

README.md

@@ -54,6 +54,22 @@ class OptILMClassifier(nn.Module):
         logits = self.classifier(combined_input)
         return logits
 
+
+def load_optillm_model():
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
+    # Load the base model
+    base_model = AutoModel.from_pretrained("google-bert/bert-large-uncased")
+    # Create the OptILMClassifier
+    model = OptILMClassifier(base_model, num_labels=len(APPROACHES))  
+    model.to(device)
+    # Download the safetensors file
+    safetensors_path = hf_hub_download(repo_id=MODEL_NAME, filename="model.safetensors")
+    # Load the state dict from the safetensors file
+    load_model(model, safetensors_path)
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    return model, tokenizer, device
+
 def preprocess_input(tokenizer, system_prompt, initial_query):
     combined_input = f"{system_prompt}\n\nUser: {initial_query}"
     encoding = tokenizer.encode_plus(