Upload app.py

app.py

@@ -1,43 +1,49 @@
-#!/usr/bin/env python3
-import gradio as gr
+import streamlit as st
 import os
 
+st.set_page_config(page_title='ML Model Trainer', page_icon='🤖')
+
+st.title('🤖 ML Model Trainer')
+st.markdown('Generate training scripts for SFT, DPO, LoRA fine-tuning')
+
 # Model options
 MODELS = [
-    # Small (fast, <3B)
     'Qwen/Qwen2.5-0.5B-Instruct',
-    'Qwen/Qwen2.5-1.5B-Instruct', 
-    'microsoft/Phi-3-mini-128k-instruct',
-    'google/gemma-2b-it',
-    'meta-llama/Llama-3.2-1B-Instruct',
-    # Medium (7-13B)
+    'Qwen/Qwen2.5-1.5B-Instruct',
     'Qwen/Qwen2.5-7B-Instruct',
+    'meta-llama/Llama-3.2-1B-Instruct',
     'meta-llama/Llama-3.2-3B-Instruct',
+    'microsoft/Phi-3-mini-128k-instruct',
+    'google/gemma-2b-it',
     'mistralai/Mistral-7B-Instruct-v0.3',
-    # Large (needs more GPU)
-    'Qwen/Qwen2.5-14B-Instruct',
-    'mistralai/Mixtral-8x7B-Instruct-v0.1',
 ]
 
-# Training methods
-METHODS = ['SFT', 'DPO', 'GRPO', 'LoRA']
+METHODS = ['SFT', 'DPO', 'LoRA']
 
-# Public datasets
 DATASETS = [
     'HuggingFaceH4/ultrachat_200k',
     'openai/gsm8k',
     'meta-math/MATH',
     'anthropic/hh-rlhf',
-    'stanfordnlp/SHP',
 ]
 
-def generate_training_script(model, method, dataset, epochs, lr, batch_size, max_length, output_name):
-    template = f'''#!/usr/bin/env python3
-# Auto-generated training script
-# Model: {model}
-# Method: {method}
-# Dataset: {dataset}
+col1, col2 = st.columns(2)
+
+with col1:
+    model = st.selectbox('Model', MODELS)
+    method = st.selectbox('Training Method', METHODS)
+    dataset = st.selectbox('Dataset', DATASETS)
+    output_name = st.text_input('Output Model Name', 'my-finetuned-model')
 
+with col2:
+    epochs = st.slider('Epochs', 1, 10, 3)
+    lr = st.number_input('Learning Rate', value=2e-5, format='%.0e')
+    batch_size = st.slider('Batch Size', 1, 16, 4)
+    max_length = st.slider('Max Sequence Length', 256, 8192, 4096, 256)
+
+if st.button('🔧 Generate Training Script', type='primary'):
+    if method == 'SFT':
+        script = f'''#!/usr/bin/env python3
 from trl import SFTTrainer, SFTConfig
 from datasets import load_dataset
 
@@ -67,14 +73,8 @@ trainer = SFTTrainer(
 trainer.train()
 trainer.push_to_hub()
 '''
-    return template
-
-def generate_dpo_script(model, dataset, epochs, lr, batch_size, output_name):
-    template = f'''#!/usr/bin/env python3
-# DPO Training Script
-# Model: {model}
-# Dataset: {dataset}
-
+    elif method == 'DPO':
+        script = f'''#!/usr/bin/env python3
 from trl import DPOTrainer, DPOConfig
 from datasets import load_dataset
 
@@ -84,7 +84,7 @@ training_args = DPOConfig(
     learning_rate={lr},
     num_train_epochs={epochs},
     per_device_train_batch_size={batch_size},
-    max_seq_length={512},
+    max_seq_length=512,
     bf16=True,
     output_dir='./{output_name}-output',
     push_to_hub=True,
@@ -100,14 +100,8 @@ trainer = DPOTrainer(
 trainer.train()
 trainer.push_to_hub()
 '''
-    return template
-
-def generate_lora_script(model, dataset, epochs, lr, batch_size, output_name):
-    template = f'''#!/usr/bin/env python3
-# LoRA Fine-tuning Script
-# Model: {model}
-# Dataset: {dataset}
-
+    else:  # LoRA
+        script = f'''#!/usr/bin/env python3
 from trl import SFTTrainer, SFTConfig
 from peft import LoraConfig
 from datasets import load_dataset
@@ -122,7 +116,7 @@ peft_config = LoraConfig(
 )
 
 training_args = SFTConfig(
-    learning_rate={lr * 10},  # LoRA needs higher LR
+    learning_rate={lr * 10},
     num_train_epochs={epochs},
     per_device_train_batch_size={batch_size},
     max_seq_length=2048,
@@ -142,69 +136,27 @@ trainer = SFTTrainer(
 trainer.train()
 trainer.push_to_hub()
 '''
-    return template
-
-def generate_script(model, method, dataset, epochs, lr, batch_size, max_length, output_name):
-    if method == 'DPO':
-        return generate_dpo_script(model, dataset, epochs, lr, batch_size, output_name)
-    elif method == 'LoRA':
-        return generate_lora_script(model, dataset, epochs, lr, batch_size, output_name)
-    else:  # SFT, GRPO
-        return generate_training_script(model, method, dataset, epochs, lr, batch_size, max_length, output_name)
-
-def get_hardware_requirement(model):
+    
+    st.code(script, language='python')
+    
+    # Hardware info
     if '0.5B' in model or '1B' in model:
-        return 'a10g-small (24GB VRAM)'
+        hw = 'a10g-small (24GB VRAM) - ~2h training'
     elif '3B' in model or '7B' in model:
-        return 'a10g-large or a100-small (24-80GB VRAM)'
+        hw = 'a10g-large or a100-small (24-80GB VRAM) - ~4h training'
     else:
-        return 'a100-large (80GB VRAM)'
-
-# Gradio UI
-with gr.Blocks(title='ML Model Trainer', theme=gr.themes.Soft()) as demo:
-    gr.Markdown('# 🤖 ML Model Trainer')
-    gr.Markdown('Generate training scripts for SFT, DPO, LoRA fine-tuning')
-    
-    with gr.Row():
-        with gr.Column():
-            model = gr.Dropdown(MODELS, label='Model', value='Qwen/Qwen2.5-0.5B-Instruct')
-            method = gr.Dropdown(METHODS, label='Training Method', value='SFT')
-            dataset = gr.Dropdown(DATASETS, label='Dataset', value='HuggingFaceH4/ultrachat_200k')
-            output_name = gr.Textbox(label='Output Model Name', value='my-finetuned-model')
-        
-        with gr.Column():
-            epochs = gr.Slider(1, 10, value=3, step=1, label='Epochs')
-            lr = gr.Number(label='Learning Rate', value=2e-5)
-            batch_size = gr.Slider(1, 16, value=4, step=1, label='Batch Size')
-            max_length = gr.Slider(256, 8192, value=4096, step=256, label='Max Sequence Length')
-    
-    generate_btn = gr.Button('🔧 Generate Training Script', variant='primary')
+        hw = 'a100-large (80GB VRAM) - ~6h training'
     
-    output_code = gr.Code(label='Training Script', language='python', lines=20)
-    
-    hardware_info = gr.Markdown('')
-    
-    def on_generate(model, method, dataset, epochs, lr, batch_size, max_length, output_name):
-        script = generate_script(model, method, dataset, epochs, lr, batch_size, max_length, output_name)
-        hw = get_hardware_requirement(model)
-        return script, f'**Recommended Hardware:** {hw} | **Timeout:** ~{int(epochs * 2)}h'
-    
-    generate_btn.click(
-        on_generate,
-        inputs=[model, method, dataset, epochs, lr, batch_size, max_length, output_name],
-        outputs=[output_code, hardware_info]
-    )
-    
-    gr.Markdown('---')
-    gr.Markdown('### 📋 How to Use')
-    gr.Markdown('''
-    1. Configure your training parameters above
-    2. Click **Generate Training Script**
-    3. Copy the script to a file (e.g., `train.py`)
-    4. Install dependencies: `pip install transformers trl torch datasets accelerate peft`
-    5. Run: `python train.py`
-    
-    **Note:** You'll need Hugging Face Pro or compute credits for cloud training.
-    ''')
-
-demo.launch()
+    st.info(f'**Recommended:** {hw}')
+
+st.markdown('---')
+st.markdown('### 📋 How to Use')
+st.markdown('''
+1. Configure parameters above
+2. Click **Generate Training Script**
+3. Copy the script to `train.py`
+4. Install: `pip install transformers trl torch datasets accelerate peft`
+5. Run: `python train.py`
+
+**Note:** Need Hugging Face Pro or compute credits for cloud training.
+''')