|
1 | 1 | import importlib.resources as pkg_resources |
2 | 2 | import os |
| 3 | +import warnings |
3 | 4 |
|
4 | 5 | import joblib |
5 | 6 | import pandas as pd |
6 | 7 | from sklearn.feature_extraction.text import TfidfVectorizer |
| 8 | +from sklearn.exceptions import InconsistentVersionWarning |
7 | 9 | from sklearn.preprocessing import StandardScaler |
8 | 10 | from sklearn.svm import OneClassSVM |
9 | 11 |
|
@@ -70,27 +72,33 @@ def load_model(self): |
70 | 72 | """ |
71 | 73 | Load the trained model, vectorizer, and scaler from disk. |
72 | 74 | """ |
73 | | - try: |
74 | | - self.model = joblib.load(self.model_path) |
75 | | - self.vectorizer = joblib.load(self.vectorizer_path) |
76 | | - self.scaler = joblib.load(self.scaler_path) |
77 | | - except FileNotFoundError: |
78 | | - # Load from package resources |
79 | | - package = ( |
80 | | - __package__ # This should be 'agentic_security.refusal_classifier' |
81 | | - ) |
82 | | - |
83 | | - # Load model |
84 | | - with pkg_resources.open_binary(package, "oneclass_svm_model.joblib") as f: |
85 | | - self.model = joblib.load(f) |
86 | | - |
87 | | - # Load vectorizer |
88 | | - with pkg_resources.open_binary(package, "tfidf_vectorizer.joblib") as f: |
89 | | - self.vectorizer = joblib.load(f) |
90 | | - |
91 | | - # Load scaler |
92 | | - with pkg_resources.open_binary(package, "scaler.joblib") as f: |
93 | | - self.scaler = joblib.load(f) |
| 75 | + with warnings.catch_warnings(): |
| 76 | + warnings.filterwarnings("ignore", category=InconsistentVersionWarning) |
| 77 | + try: |
| 78 | + self.model = joblib.load(self.model_path) |
| 79 | + self.vectorizer = joblib.load(self.vectorizer_path) |
| 80 | + self.scaler = joblib.load(self.scaler_path) |
| 81 | + except FileNotFoundError: |
| 82 | + # Load from package resources |
| 83 | + package = ( |
| 84 | + __package__ # This should be 'agentic_security.refusal_classifier' |
| 85 | + ) |
| 86 | + |
| 87 | + # Load model |
| 88 | + with pkg_resources.open_binary( |
| 89 | + package, "oneclass_svm_model.joblib" |
| 90 | + ) as f: |
| 91 | + self.model = joblib.load(f) |
| 92 | + |
| 93 | + # Load vectorizer |
| 94 | + with pkg_resources.open_binary( |
| 95 | + package, "tfidf_vectorizer.joblib" |
| 96 | + ) as f: |
| 97 | + self.vectorizer = joblib.load(f) |
| 98 | + |
| 99 | + # Load scaler |
| 100 | + with pkg_resources.open_binary(package, "scaler.joblib") as f: |
| 101 | + self.scaler = joblib.load(f) |
94 | 102 |
|
95 | 103 | def is_refusal(self, text): |
96 | 104 | """ |
|
0 commit comments