whackthejacker/DataHubHub
1
1import streamlit as st2import os3import pandas as pd4import numpy as np5import plotly.express as px6import json7from pathlib import Path8 9# Make sure necessary directories exist10os.makedirs('assets', exist_ok=True)11os.makedirs('database/data', exist_ok=True)12os.makedirs('fine_tuned_models', exist_ok=True)13 14# Page configuration15st.set_page_config(16 page_title="ML Dataset & Code Generation Manager",17 page_icon="🤗",18 layout="wide",19 initial_sidebar_state="expanded",20)21 22def load_css():23 """Load custom CSS styles"""24 css_dir = Path("assets")25 css_path = css_dir / "custom.css"26 27 if not css_path.exists():28 # Create assets directory if it doesn't exist29 css_dir.mkdir(exist_ok=True)30 31 # Create a basic CSS file if it doesn't exist32 with open(css_path, "w") as f:33 f.write("""34 /* Custom styles for ML Dataset & Code Generation Manager */35 @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=Space+Grotesk:wght@500;700&display=swap');36 37 h1, h2, h3, h4, h5, h6 {38 font-family: 'Space Grotesk', sans-serif;39 font-weight: 700;40 color: #1A1C1F;41 }42 43 body {44 font-family: 'Inter', sans-serif;45 color: #1A1C1F;46 background-color: #F8F9FA;47 }48 49 .stButton button {50 background-color: #2563EB;51 color: white;52 border-radius: 4px;53 border: none;54 padding: 0.5rem 1rem;55 font-weight: 600;56 }57 58 .stButton button:hover {59 background-color: #1D4ED8;60 }61 62 /* Card styling */63 .card {64 background-color: white;65 border-radius: 8px;66 padding: 1.5rem;67 box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);68 margin-bottom: 1rem;69 }70 71 /* Accent colors */72 .accent-primary {73 color: #2563EB;74 }75 76 .accent-secondary {77 color: #84919A;78 }79 80 .accent-success {81 color: #10B981;82 }83 84 .accent-warning {85 color: #F59E0B;86 }87 88 .accent-danger {89 color: #EF4444;90 }91 """)92 93 # Load custom CSS94 with open(css_path, "r") as f:95 st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)96 97def render_finetune_ui():98 """99 Renders the fine-tuning UI for code generation models.100 """101 try:102 from components.fine_tuning.finetune_ui import render_finetune_ui as ft_ui103 ft_ui()104 except ImportError as e:105 st.error(f"Could not load fine-tuning UI: {e}")106 107 # Create default fine-tuning UI component if not exists108 os.makedirs("components/fine_tuning", exist_ok=True)109 if not os.path.exists("components/fine_tuning/__init__.py"):110 with open("components/fine_tuning/__init__.py", "w") as f:111 f.write('"""\nFine-tuning package for code generation models.\n"""\n')112 113 if not os.path.exists("components/fine_tuning/finetune_ui.py"):114 with open("components/fine_tuning/finetune_ui.py", "w") as f:115 f.write('''"""116Streamlit UI for fine-tuning code generation models.117"""118import streamlit as st119import pandas as pd120import os121 122def render_dataset_preparation():123 """124 Render the dataset preparation interface.125 """126 st.subheader("Dataset Preparation")127 st.write("Prepare your dataset for fine-tuning code generation models.")128 129 # Dataset upload130 uploaded_file = st.file_uploader("Upload your dataset", type=["csv", "json"])131 if uploaded_file is not None:132 try:133 if uploaded_file.name.endswith('.csv'):134 df = pd.read_csv(uploaded_file)135 else:136 df = pd.read_json(uploaded_file)137 138 st.write("Dataset Preview:")139 st.dataframe(df.head())140 141 # Example of data columns mapping142 st.subheader("Column Mapping")143 144 input_col = st.selectbox("Select input column (e.g., code)", df.columns)145 target_col = st.selectbox("Select target column (e.g., comment)", df.columns)146 147 # Sample transformation148 if st.button("Apply Transformation"):149 if input_col and target_col:150 # Example transformation: simple trim/clean151 df[input_col] = df[input_col].astype(str).str.strip()152 df[target_col] = df[target_col].astype(str).str.strip()153 154 st.write("Transformed Dataset:")155 st.dataframe(df.head())156 157 # Option to save processed dataset158 if st.button("Save Processed Dataset"):159 processed_path = os.path.join("datasets", "processed_dataset.csv")160 os.makedirs("datasets", exist_ok=True)161 df.to_csv(processed_path, index=False)162 st.success(f"Dataset saved to {processed_path}")163 except Exception as e:164 st.error(f"Error processing dataset: {e}")165 166def render_model_training():167 """168 Render the model training interface.169 """170 st.subheader("Model Training")171 st.write("Configure and start training your model.")172 173 # Model selection174 model_options = [175 "Salesforce/codet5-small",176 "Salesforce/codet5-base",177 "microsoft/codebert-base",178 "microsoft/graphcodebert-base"179 ]180 181 selected_model = st.selectbox("Select base model", model_options)182 183 # Training parameters184 col1, col2 = st.columns(2)185 with col1:186 batch_size = st.number_input("Batch size", min_value=1, max_value=64, value=8)187 epochs = st.number_input("Number of epochs", min_value=1, max_value=100, value=3)188 learning_rate = st.number_input("Learning rate", min_value=0.00001, max_value=0.1, value=0.0001, format="%.5f")189 190 with col2:191 max_input_length = st.number_input("Max input length", min_value=32, max_value=512, value=128)192 max_target_length = st.number_input("Max target length", min_value=32, max_value=512, value=128)193 task_type = st.selectbox("Task type", ["Code to Comment", "Comment to Code"])194 195 # Training button (placeholder)196 if st.button("Start Training"):197 st.info("Training would start here. This is a placeholder.")198 # In a real implementation, this would call the training function199 # and display a progress bar or redirect to a training monitoring page200 201def render_model_testing():202 """203 Render the model testing interface.204 """205 st.subheader("Model Testing")206 st.write("Test your fine-tuned model with custom inputs.")207 208 # Model selection209 st.selectbox("Select fine-tuned model", ["No models available yet"])210 211 # Test input212 if st.selectbox("Task type", ["Code to Comment", "Comment to Code"]) == "Code to Comment":213 test_input = st.text_area("Enter code to generate a comment", 214 value="def fibonacci(n):\\n if n <= 1:\\n return n\\n else:\\n return fibonacci(n-1) + fibonacci(n-2)")215 placeholder = "# This function implements the Fibonacci sequence recursively..."216 else:217 test_input = st.text_area("Enter comment to generate code", 218 value="# A function that calculates the factorial of a number recursively")219 placeholder = "def factorial(n):\\n if n == 0:\\n return 1\\n else:\\n return n * factorial(n-1)"220 221 # Generate button (placeholder)222 if st.button("Generate"):223 st.code(placeholder, language="python")224 # In a real implementation, this would call the model inference function225 226def render_finetune_ui():227 """228 Render the fine-tuning UI for code generation models.229 """230 st.title("Fine-Tune Code Generation Models")231 232 tabs = st.tabs(["Dataset Preparation", "Model Training", "Model Testing"])233 234 with tabs[0]:235 render_dataset_preparation()236 237 with tabs[1]:238 render_model_training()239 240 with tabs[2]:241 render_model_testing()242''')243 244 # Try again after creating the files245 try:246 from components.fine_tuning.finetune_ui import render_finetune_ui as ft_ui247 ft_ui()248 except ImportError as e:249 st.error(f"Still could not load fine-tuning UI after creating files: {e}")250 st.info("Please restart the app to initialize the components.")251 252def render_code_quality_ui():253 """254 Renders the code quality tools UI.255 """256 try:257 from components.code_quality import render_code_quality_tools258 render_code_quality_tools()259 except ImportError:260 st.error("Code quality tools not found. Implementing basic version.")261 st.title("Code Quality Tools")262 st.write("This section will provide tools for code linting, formatting, and testing.")263 264 # Tabs for different code quality tools265 tabs = st.tabs(["Linting", "Formatting", "Type Checking", "Testing"])266 267 with tabs[0]:268 st.subheader("Code Linting")269 st.write("Tools for checking code quality and style.")270 st.code("# Coming soon: PyLint and Flake8 integration")271 272 with tabs[1]:273 st.subheader("Code Formatting")274 st.write("Tools for formatting code according to style guides.")275 st.code("# Coming soon: Black and isort integration")276 277 with tabs[2]:278 st.subheader("Type Checking")279 st.write("Tools for checking type annotations.")280 st.code("# Coming soon: MyPy integration")281 282 with tabs[3]:283 st.subheader("Testing")284 st.write("Tools for running tests and checking code coverage.")285 st.code("# Coming soon: PyTest integration")286 287def render_dataset_management_ui():288 """289 Renders the dataset management UI.290 """291 st.title("Dataset Management")292 293 # Tabs for different dataset operations294 tabs = st.tabs(["Upload", "Preview", "Statistics", "Visualization", "Validation", "Version Control"])295 296 with tabs[0]:297 try:298 from components.dataset_uploader import render_dataset_uploader299 render_dataset_uploader()300 except ImportError:301 st.subheader("Dataset Upload")302 st.write("Upload your datasets in CSV or JSON format.")303 304 uploaded_file = st.file_uploader("Choose a file", type=["csv", "json"])305 if uploaded_file is not None:306 try:307 if uploaded_file.name.endswith('.csv'):308 df = pd.read_csv(uploaded_file)309 dataset_type = "csv"310 else:311 df = pd.read_json(uploaded_file)312 dataset_type = "json"313 314 st.session_state["dataset"] = df315 st.session_state["dataset_type"] = dataset_type316 st.success(f"Successfully loaded {dataset_type.upper()} file with {df.shape[0]} rows and {df.shape[1]} columns.")317 st.dataframe(df.head())318 except Exception as e:319 st.error(f"Error: {e}")320 321 with tabs[1]:322 if "dataset" in st.session_state:323 try:324 from components.dataset_preview import render_dataset_preview325 render_dataset_preview(st.session_state["dataset"], st.session_state["dataset_type"])326 except ImportError:327 st.subheader("Dataset Preview")328 st.dataframe(st.session_state["dataset"].head(10))329 else:330 st.info("Please upload a dataset first.")331 332 with tabs[2]:333 if "dataset" in st.session_state:334 try:335 from components.dataset_statistics import render_dataset_statistics336 render_dataset_statistics(st.session_state["dataset"], st.session_state["dataset_type"])337 except ImportError:338 st.subheader("Dataset Statistics")339 st.write("Basic statistics:")340 st.write(st.session_state["dataset"].describe())341 342 # Missing values343 missing_data = st.session_state["dataset"].isnull().sum()344 st.write("Missing values per column:")345 st.write(missing_data[missing_data > 0])346 else:347 st.info("Please upload a dataset first.")348 349 with tabs[3]:350 if "dataset" in st.session_state:351 try:352 from components.dataset_visualization import render_dataset_visualization353 render_dataset_visualization(st.session_state["dataset"], st.session_state["dataset_type"])354 except ImportError:355 st.subheader("Dataset Visualization")356 357 # Only show for numerical columns358 numeric_cols = st.session_state["dataset"].select_dtypes(include=[np.number]).columns.tolist()359 360 if len(numeric_cols) > 0:361 col1, col2 = st.columns(2)362 363 with col1:364 x_axis = st.selectbox("X-axis", numeric_cols)365 366 with col2:367 y_axis = st.selectbox("Y-axis", numeric_cols, index=min(1, len(numeric_cols)-1))368 369 fig = px.scatter(st.session_state["dataset"], x=x_axis, y=y_axis)370 st.plotly_chart(fig, use_container_width=True)371 else:372 st.write("No numerical columns available for visualization.")373 else:374 st.info("Please upload a dataset first.")375 376 with tabs[4]:377 if "dataset" in st.session_state:378 try:379 from components.dataset_validation import render_dataset_validation380 render_dataset_validation(st.session_state["dataset"], st.session_state["dataset_type"])381 except ImportError:382 st.subheader("Dataset Validation")383 384 # Simple validation checks385 st.write("Dataset Shape:", st.session_state["dataset"].shape)386 st.write("Duplicate Rows:", st.session_state["dataset"].duplicated().sum())387 388 # Missing values percentage389 missing_percent = (st.session_state["dataset"].isnull().sum() / len(st.session_state["dataset"])) * 100390 st.write("Missing Values Percentage:")391 st.write(missing_percent[missing_percent > 0])392 else:393 st.info("Please upload a dataset first.")394 395 with tabs[5]:396 if "dataset" in st.session_state:397 try:398 from components.dataset_version_control import render_version_control_ui, render_save_version_ui, render_version_visualization399 400 # If we have a dataset ID in session state, use it, otherwise prompt to save first401 if "dataset_id" in st.session_state:402 dataset_id = st.session_state["dataset_id"]403 404 # Show dataset version control UI405 render_version_control_ui(dataset_id, st.session_state.get("dataset"))406 407 # Show save version UI408 st.divider()409 if st.session_state.get("dataset") is not None:410 new_version = render_save_version_ui(dataset_id, st.session_state["dataset"])411 if new_version:412 st.success(f"Created new version: {new_version.version_id}")413 414 # Show version visualization415 st.divider()416 render_version_visualization(dataset_id)417 else:418 # No dataset ID yet, so prompt to save the dataset first419 st.info("To use version control, first save this dataset to the database.")420 421 dataset_name = st.text_input("Dataset Name", value="My Dataset")422 dataset_description = st.text_area("Dataset Description", value="Dataset uploaded for analysis")423 424 if st.button("Save Dataset to Database"):425 # Import database operations426 from database.operations import DatasetOperations, DatasetVersionOperations427 428 # Store dataset in database429 dataset = DatasetOperations.store_dataframe_info(430 df=st.session_state["dataset"],431 name=dataset_name,432 description=dataset_description,433 source="local_upload"434 )435 436 # Store as initial version437 initial_version = DatasetVersionOperations.create_version_from_dataframe(438 dataset_id=dataset.id,439 df=st.session_state["dataset"],440 description="Initial version"441 )442 443 # Store dataset ID in session state444 st.session_state["dataset_id"] = dataset.id445 446 st.success(f"Dataset saved to database with ID: {dataset.id}")447 st.success(f"Initial version created: {initial_version.version_id}")448 449 # Rerun to show version control UI450 st.experimental_rerun()451 except ImportError as e:452 st.subheader("Dataset Version Control")453 st.error(f"Could not load version control components: {e}")454 st.info("Please make sure all required components are installed.")455 else:456 st.info("Please upload a dataset first.")457 458def main():459 """460 Main function to run the application.461 """462 # Load custom CSS463 load_css()464 465 # Sidebar for navigation466 st.sidebar.title("ML Dataset & Code Gen Manager")467 468 # Navigation469 page = st.sidebar.radio("Navigation", ["Home", "Dataset Management", "Fine-Tuning", "Code Quality Tools"])470 471 # Display selected page472 if page == "Home":473 st.title("ML Dataset & Code Generation Manager")474 st.write("Welcome to the ML Dataset & Code Generation Manager. This platform helps you manage ML datasets and fine-tune code generation models.")475 476 # Main features in cards477 col1, col2 = st.columns(2)478 479 with col1:480 st.markdown("""481 <div class="card">482 <h3>Dataset Management</h3>483 <p>Upload, analyze, visualize, and validate your ML datasets.</p>484 <ul>485 <li>Support for CSV and JSON formats</li>486 <li>Statistical analysis and visualization</li>487 <li>Data validation and quality checks</li>488 <li>Hugging Face Hub integration</li>489 </ul>490 </div>491 """, unsafe_allow_html=True)492 493 st.markdown("""494 <div class="card">495 <h3>Code Quality Tools</h3>496 <p>Tools for ensuring high-quality code.</p>497 <ul>498 <li>Code linting with PyLint</li>499 <li>Code formatting with Black and isort</li>500 <li>Type checking with MyPy</li>501 <li>Testing with PyTest</li>502 </ul>503 </div>504 """, unsafe_allow_html=True)505 506 with col2:507 st.markdown("""508 <div class="card">509 <h3>Fine-Tuning</h3>510 <p>Fine-tune code generation models on your custom datasets.</p>511 <ul>512 <li>Support for CodeT5, CodeBERT models</li>513 <li>Code-to-comment and comment-to-code tasks</li>514 <li>Custom dataset preparation</li>515 <li>Model testing and evaluation</li>516 </ul>517 </div>518 """, unsafe_allow_html=True)519 520 st.markdown("""521 <div class="card">522 <h3>Hugging Face Integration</h3>523 <p>Seamless integration with Hugging Face Hub.</p>524 <ul>525 <li>Search and load models and datasets</li>526 <li>Deploy fine-tuned models to Hugging Face Spaces</li>527 <li>Share and collaborate on models and datasets</li>528 </ul>529 </div>530 """, unsafe_allow_html=True)531 532 # Get started section533 st.subheader("Get Started")534 st.write("To get started, navigate to the Dataset Management page to upload your data, or explore the Fine-Tuning page to train code generation models.")535 536 elif page == "Dataset Management":537 render_dataset_management_ui()538 539 elif page == "Fine-Tuning":540 render_finetune_ui()541 542 elif page == "Code Quality Tools":543 render_code_quality_ui()544 545if __name__ == "__main__":546 main()