wasdcutecat's picture
Update app.py
bf5a6c9 verified
import gradio as gr
import pickle
import numpy as np
# Modeli yükle
with open("xgboost_optimized_v2.pkl", "rb") as f: # model/ klasörü kaldırıldı
model_pkg = pickle.load(f)
# Modeller ve fonksiyonlar
xgb_improved = model_pkg["base_models"]["xgboost_improved"][0]
smart_align = model_pkg["enhancement_functions"]["smart_feature_alignment"]
ensemble_fn = model_pkg["enhancement_functions"]["ensemble_predict"]
unknown_boost_fn = model_pkg["enhancement_functions"]["unknown_category_boost"]
dynamic_adjust_fn = model_pkg["enhancement_functions"]["dynamic_risk_adjustment"]
confidence_fn = model_pkg["enhancement_functions"]["calculate_confidence"]
# Ön işleme bileşenleri
preprocessors = model_pkg["preprocessors"]
scaler = preprocessors["scaler"]
encoders = preprocessors["encoders"]
bert_extractor = preprocessors["bert_extractor"]
label_encoder = preprocessors["label_encoder"]
feature_names = scaler.mean_.shape[0] # Özellik sayısı
# Tahmin fonksiyonu
def predict(cve_id, cwe_id, description, cvss_score):
# Dummy girişler (GPU olmadığı için BERT çıkarımı yapılmıyor)
numeric = np.zeros((1, feature_names))
bert_feats = np.zeros((1, 768)) # Dummy BERT vektörü
x_input = np.hstack([numeric, bert_feats])
# Model tahminleri
xgb_pred = xgb_improved.predict(x_input)[0]
ensemble_pred = ensemble_fn(cvss_score, xgb_pred)
final_pred = unknown_boost_fn(ensemble_pred, unknown_features=2, cvss_score=cvss_score)
optimized = dynamic_adjust_fn(final_pred, cvss_score, cwe_id, len(description))
confidence = confidence_fn(2, cvss_score)
# Kategori sınıflandırması
category = "Düşük" if optimized <= 3 else "Orta" if optimized <= 7 else "Yüksek"
return {
"Tahmini Risk Skoru": round(optimized, 2),
"Kategori": category,
"Model Güveni": f"%{round(confidence * 100, 1)}"
}
# Gradio arayüzü
interface = gr.Interface(
fn=predict,
inputs=[
gr.Text(label="CVE ID"),
gr.Text(label="CWE ID"),
gr.Textbox(label="Açıklama"),
gr.Slider(minimum=0.0, maximum=10.0, step=0.1, label="CVSS Skoru")
],
outputs="json",
title="Risk Skoru Tahmin Modeli (XGBoost Optimized)",
description="CVE tanımlarına göre tahmini risk skoru ve kategori"
)
# Uygulamayı başlat
interface.launch()