diff --git a/models/decision_tree.py b/models/decision_tree.py new file mode 100644 index 0000000..c790a0c --- /dev/null +++ b/models/decision_tree.py @@ -0,0 +1,37 @@ +from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor +from sklearn.metrics import accuracy_score, mean_squared_error +import pandas as pd +import numpy as np + +class DecisionTreeModel: + """ + Decision Tree model wrapper supporting both classification and regression. + """ + + def __init__(self, task_type="classification", max_depth=None, random_state=42): + self.task_type = task_type + self.max_depth = max_depth + self.random_state = random_state + + if task_type == "classification": + self.model = DecisionTreeClassifier(max_depth=max_depth, random_state=random_state) + elif task_type == "regression": + self.model = DecisionTreeRegressor(max_depth=max_depth, random_state=random_state) + else: + raise ValueError("Invalid task_type. Use 'classification' or 'regression'.") + + def train(self, X, y): + self.model.fit(X, y) + + def predict(self, X): + return self.model.predict(X) + + def evaluate(self, X, y): + preds = self.predict(X) + if self.task_type == "classification": + return {"accuracy": accuracy_score(y, preds)} + else: + return {"mse": mean_squared_error(y, preds)} + + def get_model(self): + return self.model diff --git a/pages/Decision_Tree.py b/pages/Decision_Tree.py new file mode 100644 index 0000000..ab10c27 --- /dev/null +++ b/pages/Decision_Tree.py @@ -0,0 +1,52 @@ +import streamlit as st +import pandas as pd +from models.decision_tree import DecisionTreeModel +from utils.data_helpers import generate_sample_regression_data # For regression data +from sklearn.datasets import make_classification +import matplotlib.pyplot as plt +from sklearn.tree import plot_tree + +st.title("🌳 Decision Tree Model Playground") + +task_type = st.radio("Select Task Type", ["classification", "regression"]) + +# Sidebar parameters +st.sidebar.header("Model Parameters") +max_depth = st.sidebar.slider("Max Depth", 1, 10, 3) +n_samples = st.sidebar.slider("Number of Samples", 50, 500, 100) +n_features = st.sidebar.slider("Number of Features", 1, 5, 2) +noise = st.sidebar.slider("Noise (for regression)", 0.0, 10.0, 1.0) + +st.write("### Generated Dataset Preview") + +# Generate dataset +if task_type == "classification": + X, y = make_classification( + n_samples=n_samples, n_features=n_features, n_informative=n_features, + n_redundant=0, random_state=42 + ) + df = pd.DataFrame(X, columns=[f"X{i+1}" for i in range(n_features)]) + df["y"] = y +else: + df = generate_sample_regression_data(n_samples, n_features, noise) + +st.dataframe(df.head()) + +# Train model +st.write("### Train Model") +model = DecisionTreeModel(task_type=task_type, max_depth=max_depth) +X = df.iloc[:, :-1] +y = df.iloc[:, -1] +model.train(X, y) + +# Evaluate +metrics = model.evaluate(X, y) +st.write("### Evaluation Metrics") +st.json(metrics) + +# Visualization +st.write("### Decision Tree Visualization") + +fig, ax = plt.subplots(figsize=(10, 6)) +plot_tree(model.get_model(), filled=True, feature_names=X.columns) +st.pyplot(fig)