Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import pandas as pd | |
| import numpy as np | |
| import tensorflow as tf | |
| from sklearn.preprocessing import MinMaxScaler | |
| import plotly.express as px | |
| import os | |
| from groq import Groq | |
| # Initialize session state | |
| if 'model' not in st.session_state: | |
| st.session_state.model = None | |
| if 'threshold' not in st.session_state: | |
| st.session_state.threshold = None | |
| # Autoencoder model definition | |
| def build_autoencoder(input_dim): | |
| model = tf.keras.Sequential([ | |
| tf.keras.layers.Dense(64, activation='relu', input_shape=(input_dim,)), | |
| tf.keras.layers.Dense(32, activation='relu'), | |
| tf.keras.layers.Dense(64, activation='relu'), | |
| tf.keras.layers.Dense(input_dim) | |
| ]) | |
| model.compile(optimizer='adam', loss='mse') | |
| return model | |
| # Sidebar controls | |
| st.sidebar.title("Configuration") | |
| fine_tune = st.sidebar.button("Fine-tune Model") | |
| groq_api_key = st.sidebar.text_input("Groq API Key (optional)", type="password") | |
| # Main interface | |
| st.title("🛰️ AI Network Anomaly Detector") | |
| st.write("Upload your network data (CSV) to detect anomalies") | |
| # File uploader | |
| uploaded_file = st.file_uploader("Choose a CSV file", type="csv") | |
| # Load or generate sample data | |
| if uploaded_file is not None: | |
| data = pd.read_csv(uploaded_file) | |
| else: | |
| st.info("Using sample data. Upload a file to use your own dataset.") | |
| data = pd.read_csv("sample_wifi_data.csv") # You should provide this sample file | |
| # Preprocessing | |
| features = ['device_count', 'connection_attempts', 'packet_loss', 'latency'] | |
| scaler = MinMaxScaler() | |
| data_scaled = scaler.fit_transform(data[features]) | |
| # Model training/fine-tuning | |
| if fine_tune or st.session_state.model is None: | |
| with st.spinner("Training model..."): | |
| autoencoder = build_autoencoder(data_scaled.shape[1]) | |
| autoencoder.fit(data_scaled, data_scaled, | |
| epochs=100, | |
| batch_size=32, | |
| verbose=0, | |
| validation_split=0.1) | |
| st.session_state.model = autoencoder | |
| # Calculate threshold | |
| reconstructions = autoencoder.predict(data_scaled) | |
| mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1) | |
| st.session_state.threshold = np.percentile(mse, 95) | |
| # Anomaly detection | |
| if st.session_state.model and st.button("Detect Anomalies"): | |
| reconstructions = st.session_state.model.predict(data_scaled) | |
| mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1) | |
| anomalies = mse > st.session_state.threshold | |
| # Visualization | |
| fig = px.line(data, x=data.index, y='connection_attempts', | |
| title='Network Traffic with Anomalies') | |
| fig.add_scatter(x=data[anomalies].index, y=data[anomalies]['connection_attempts'], | |
| mode='markers', name='Anomalies') | |
| st.plotly_chart(fig) | |
| # Generate alert with Groq/Llama3 | |
| if groq_api_key: | |
| try: | |
| client = Groq(api_key=groq_api_key) | |
| response = client.chat.completions.create( | |
| model="llama3-70b-8192", | |
| messages=[{ | |
| "role": "user", | |
| "content": f"Generate a network security alert for {sum(anomalies)} anomalies detected. Max connection attempts: {data['connection_attempts'].max()}" | |
| }] | |
| ) | |
| st.warning(response.choices[0].message.content) | |
| except Exception as e: | |
| st.error(f"Groq API Error: {str(e)}") | |
| else: | |
| st.warning(f"Detected {sum(anomalies)} anomalies! Consider adding Groq API key for detailed analysis.") | |
| # Download button for results | |
| if st.session_state.threshold: | |
| st.download_button( | |
| label="Download Anomaly Report", | |
| data=data[anomalies].to_csv().encode('utf-8'), | |
| file_name='anomalies_report.csv', | |
| mime='text/csv' | |
| ) |