diff --git a/.flake8 b/.flake8 index c0bc336..b71e1f9 100644 --- a/.flake8 +++ b/.flake8 @@ -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 diff --git a/agentic_security/refusal_classifier/model copy.py b/agentic_security/refusal_classifier/model copy.py deleted file mode 100644 index a1db45c..0000000 --- a/agentic_security/refusal_classifier/model copy.py +++ /dev/null @@ -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 diff --git a/agentic_security/refusal_classifier/model.py b/agentic_security/refusal_classifier/model.py index 1211ddc..b217f31 100644 --- a/agentic_security/refusal_classifier/model.py +++ b/agentic_security/refusal_classifier/model.py @@ -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: diff --git a/pyproject.toml b/pyproject.toml index ca3adf5..57d8ac7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 "] maintainers = ["Alexander Miasoiedov "]