Compare commits

...

3 Commits

Author SHA1 Message Date
Alexander Myasoedov d7f6c7bd30 fix(pkg_resources.open_binary): 2024-10-19 16:31:08 +03:00
Alexander Myasoedov 6759cb0acc feat(add py3.12): 2024-10-19 16:18:58 +03:00
Alexander Myasoedov 0ab314c367 fix(module): 2024-10-19 16:17:39 +03:00
5 changed files with 29 additions and 49 deletions
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+23 -3
View File
@@ -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.
""" """
try:
self.model = joblib.load(self.model_path) self.model = joblib.load(self.model_path)
self.vectorizer = joblib.load(self.vectorizer_path) self.vectorizer = joblib.load(self.vectorizer_path)
self.scaler = joblib.load(self.scaler_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.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>"]