forked from SahrJohn/RSI-Stoch-MA-EA
31 lines
904 B
Python
31 lines
904 B
Python
|
import ast
|
||
|
|
import pandas as pd
|
||
|
|
import numpy as np
|
||
|
|
from sklearn.model_selection import train_test_split
|
||
|
|
from sklearn.metrics import roc_auc_score
|
||
|
|
from sklearn.ensemble import GradientBoostingClassifier
|
||
|
|
from sklearn.calibration import CalibratedClassifierCV
|
||
|
|
import joblib
|
||
|
|
|
||
|
|
CSV = "nn_dataset_signalready.csv"
|
||
|
|
OUT = "qt_quality_model.joblib"
|
||
|
|
|
||
|
|
df = pd.read_csv(CSV)
|
||
|
|
|
||
|
|
X = np.vstack(df["features_json"].apply(ast.literal_eval).values)
|
||
|
|
y = df["tp_first"].astype(int).values
|
||
|
|
|
||
|
|
X_train, X_val, y_train, y_val = train_test_split(
|
||
|
|
X, y, test_size=0.2, random_state=42, stratify=y
|
||
|
|
)
|
||
|
|
|
||
|
|
base = GradientBoostingClassifier(random_state=42)
|
||
|
|
cal = CalibratedClassifierCV(base, method="isotonic", cv=3)
|
||
|
|
cal.fit(X_train, y_train)
|
||
|
|
|
||
|
|
p = cal.predict_proba(X_val)[:, 1]
|
||
|
|
print("AUC:", roc_auc_score(y_val, p))
|
||
|
|
print("Mean P(tp_first):", p.mean())
|
||
|
|
|
||
|
|
joblib.dump({"model": cal, "n_features": X.shape[1]}, OUT)
|
||
|
|
print("Saved:", OUT)
|