mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-24 22:29:56 +02:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| d7f6c7bd30 |
@@ -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
@@ -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>"]
|
||||||
|
|||||||
Reference in New Issue
Block a user