mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-25 14:49:57 +02:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| d7f6c7bd30 | |||
| 6759cb0acc | |||
| 0ab314c367 |
@@ -2,4 +2,4 @@
|
|||||||
max-line-length = 160
|
max-line-length = 160
|
||||||
per-file-ignores =
|
per-file-ignores =
|
||||||
# Ignore docstring lints for tests
|
# Ignore docstring lints for tests
|
||||||
*: D100, D101, D102, D103, D104, D107, D105, D202, D205, D400, E501, D401
|
*: D100, D101, D102, D103, D104, D107, D105, D202, D205, D400, E501, D401, D200
|
||||||
|
|||||||
@@ -16,9 +16,9 @@ jobs:
|
|||||||
strategy:
|
strategy:
|
||||||
matrix:
|
matrix:
|
||||||
python-version:
|
python-version:
|
||||||
- "3.9"
|
|
||||||
- "3.10"
|
- "3.10"
|
||||||
- "3.11"
|
- "3.11"
|
||||||
|
- "3.12"
|
||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v3
|
- uses: actions/checkout@v3
|
||||||
- name: Install poetry
|
- name: Install poetry
|
||||||
|
|||||||
@@ -1,40 +0,0 @@
|
|||||||
import pandas as pd
|
|
||||||
from os import path
|
|
||||||
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
||||||
from sklearn.svm import OneClassSVM
|
|
||||||
from sklearn.preprocessing import StandardScaler
|
|
||||||
import joblib
|
|
||||||
|
|
||||||
# **Training and Saving**
|
|
||||||
|
|
||||||
# Load your data
|
|
||||||
df = pd.read_csv(path.expanduser("~/Downloads/data_en.csv"))
|
|
||||||
texts = pd.concat(
|
|
||||||
[df["GPT4_response"], df["ChatGPT_response"], df["Claude_response"]],
|
|
||||||
ignore_index=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Preprocess and vectorize
|
|
||||||
vectorizer = TfidfVectorizer(max_features=1000)
|
|
||||||
X = vectorizer.fit_transform(texts)
|
|
||||||
|
|
||||||
scaler = StandardScaler(with_mean=False)
|
|
||||||
X_scaled = scaler.fit_transform(X)
|
|
||||||
model = OneClassSVM(kernel="rbf", gamma="auto", nu=0.05).fit(X_scaled)
|
|
||||||
|
|
||||||
# Save the model and vectorizer to disk
|
|
||||||
joblib.dump(model, "oneclass_svm_model.joblib")
|
|
||||||
joblib.dump(vectorizer, "tfidf_vectorizer.joblib")
|
|
||||||
|
|
||||||
# **Loading and Predicting**
|
|
||||||
|
|
||||||
# Load the model and vectorizer from disk
|
|
||||||
model = joblib.load("oneclass_svm_model.joblib")
|
|
||||||
vectorizer = joblib.load("tfidf_vectorizer.joblib")
|
|
||||||
|
|
||||||
|
|
||||||
def is_refusal(text):
|
|
||||||
x = vectorizer.transform([text])
|
|
||||||
x_scaled = scaler.transform(x)
|
|
||||||
prediction = model.predict(x_scaled)
|
|
||||||
return prediction[0] == 1 # Returns True if it's a refusal response
|
|
||||||
@@ -1,9 +1,11 @@
|
|||||||
|
import importlib.resources as pkg_resources
|
||||||
|
import os
|
||||||
|
|
||||||
|
import joblib
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
from sklearn.feature_extraction.text import TfidfVectorizer
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
||||||
from sklearn.svm import OneClassSVM
|
|
||||||
from sklearn.preprocessing import StandardScaler
|
from sklearn.preprocessing import StandardScaler
|
||||||
import joblib
|
from sklearn.svm import OneClassSVM
|
||||||
import os
|
|
||||||
|
|
||||||
|
|
||||||
class RefusalClassifier:
|
class RefusalClassifier:
|
||||||
@@ -68,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.3"
|
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