fix(pkg_resources.open_binary):

This commit is contained in:
Alexander Myasoedov
2024-10-19 16:31:08 +03:00
parent 6759cb0acc
commit d7f6c7bd30
2 changed files with 23 additions and 4 deletions
+22 -3
View File
@@ -1,3 +1,4 @@
import importlib.resources as pkg_resources
import os
import joblib
@@ -69,9 +70,27 @@ class RefusalClassifier:
"""
Load the trained model, vectorizer, and scaler from disk.
"""
self.model = joblib.load(self.model_path)
self.vectorizer = joblib.load(self.vectorizer_path)
self.scaler = joblib.load(self.scaler_path)
try:
self.model = joblib.load(self.model_path)
self.vectorizer = joblib.load(self.vectorizer_path)
self.scaler = joblib.load(self.scaler_path)
except FileNotFoundError:
# Load from package resources
package = (
__package__ # This should be 'agentic_security.refusal_classifier'
)
# Load model
with pkg_resources.open_binary(package, "oneclass_svm_model.joblib") as f:
self.model = joblib.load(f)
# Load vectorizer
with pkg_resources.open_binary(package, "tfidf_vectorizer.joblib") as f:
self.vectorizer = joblib.load(f)
# Load scaler
with pkg_resources.open_binary(package, "scaler.joblib") as f:
self.scaler = joblib.load(f)
def is_refusal(self, text):
"""
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "agentic_security"
version = "0.2.4"
version = "0.2.5"
description = "Agentic LLM vulnerability scanner"
authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]