fix(module):

This commit is contained in:
Alexander Myasoedov
2024-10-19 16:17:39 +03:00
parent ecaea7997c
commit 0ab314c367
4 changed files with 6 additions and 45 deletions
+1 -1
View File
@@ -2,4 +2,4 @@
max-line-length = 160
per-file-ignores =
# 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,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
+4 -3
View File
@@ -1,9 +1,10 @@
import os
import joblib
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import OneClassSVM
from sklearn.preprocessing import StandardScaler
import joblib
import os
from sklearn.svm import OneClassSVM
class RefusalClassifier:
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "agentic_security"
version = "0.2.3"
version = "0.2.4"
description = "Agentic LLM vulnerability scanner"
authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]