Compare commits

...

1 Commits

Author SHA1 Message Date
Alexander Myasoedov d7f6c7bd30 fix(pkg_resources.open_binary): 2024-10-19 16:31:08 +03:00
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 os
import joblib import joblib
@@ -69,9 +70,27 @@ class RefusalClassifier:
""" """
Load the trained model, vectorizer, and scaler from disk. Load the trained model, vectorizer, and scaler from disk.
""" """
self.model = joblib.load(self.model_path) try:
self.vectorizer = joblib.load(self.vectorizer_path) self.model = joblib.load(self.model_path)
self.scaler = joblib.load(self.scaler_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): def is_refusal(self, text):
""" """
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry] [tool.poetry]
name = "agentic_security" name = "agentic_security"
version = "0.2.4" version = "0.2.5"
description = "Agentic LLM vulnerability scanner" description = "Agentic LLM vulnerability scanner"
authors = ["Alexander Miasoiedov <msoedov@gmail.com>"] authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"] maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]