Upload main.py

main.py

@@ -0,0 +1,295 @@
+import gradio as gr
+from utils.check_dataset import validate_dataset, generate_dataset_report
+from utils.sample_dataset import generate_sample_datasets
+from utils.model import GemmaFineTuning
+
+class GemmaUI:
+    def __init__(self):
+        self.model_instance = GemmaFineTuning()
+        self.default_params = self.model_instance.default_params
+
+    def create_ui(self):
+        """Create the Gradio interface"""
+        with gr.Blocks(title="Gemma Fine-tuning UI") as app:
+            gr.Markdown("# Gemma Model Fine-tuning Interface")
+            gr.Markdown("Upload your dataset, configure parameters, and fine-tune a Gemma model")
+
+            with gr.Tabs():
+                with gr.TabItem("1. Data Upload & Preprocessing"):
+                    with gr.Row():
+                        with gr.Column():
+                            file_upload = gr.File(label="Upload Dataset")
+                            file_format = gr.Radio(
+                                ["csv", "jsonl", "text"],
+                                label="File Format",
+                                value="csv"
+                            )
+                            preprocess_button = gr.Button("Preprocess Dataset")
+                            dataset_info = gr.TextArea(label="Dataset Information", interactive=False)
+
+                with gr.TabItem("2. Model & Hyperparameters"):
+                    with gr.Row():
+                        with gr.Column():
+                            model_name = gr.Dropdown(
+                                choices=[
+                                    "google/gemma-2b",
+                                    "google/gemma-7b",
+                                    "google/gemma-2b-it",
+                                    "google/gemma-7b-it"
+                                ],
+                                value=self.default_params["model_name"],
+                                label="Model Name",
+                                info="Select a Gemma model to fine-tune"
+                            )
+                            learning_rate = gr.Slider(
+                                minimum=1e-6,
+                                maximum=5e-4,
+                                value=self.default_params["learning_rate"],
+                                label="Learning Rate",
+                                info="Learning rate for the optimizer"
+                            )
+                            batch_size = gr.Slider(
+                                minimum=1,
+                                maximum=32,
+                                step=1,
+                                value=self.default_params["batch_size"],
+                                label="Batch Size",
+                                info="Number of samples in each training batch"
+                            )
+                            epochs = gr.Slider(
+                                minimum=1,
+                                maximum=10,
+                                step=1,
+                                value=self.default_params["epochs"],
+                                label="Epochs",
+                                info="Number of training epochs"
+                            )
+
+                        with gr.Column():
+                            max_length = gr.Slider(
+                                minimum=128,
+                                maximum=2048,
+                                step=16,
+                                value=self.default_params["max_length"],
+                                label="Max Sequence Length",
+                                info="Maximum token length for inputs"
+                            )
+                            use_lora = gr.Checkbox(
+                                value=self.default_params["use_lora"],
+                                label="Use LoRA for Parameter-Efficient Fine-tuning",
+                                info="Recommended for faster training and lower memory usage"
+                            )
+                            lora_r = gr.Slider(
+                                minimum=4,
+                                maximum=64,
+                                step=4,
+                                value=self.default_params["lora_r"],
+                                label="LoRA Rank (r)",
+                                info="Rank of the LoRA update matrices",
+                                visible=lambda: use_lora.value
+                            )
+                            lora_alpha = gr.Slider(
+                                minimum=4,
+                                maximum=64,
+                                step=4,
+                                value=self.default_params["lora_alpha"],
+                                label="LoRA Alpha",
+                                info="Scaling factor for LoRA updates",
+                                visible=lambda: use_lora.value
+                            )
+                            eval_ratio = gr.Slider(
+                                minimum=0.05,
+                                maximum=0.3,
+                                step=0.05,
+                                value=self.default_params["eval_ratio"],
+                                label="Validation Split Ratio",
+                                info="Portion of data to use for validation"
+                            )
+
+                with gr.TabItem("3. Training"):
+                    with gr.Row():
+                        with gr.Column():
+                            start_training_button = gr.Button("Start Fine-tuning")
+                            stop_training_button = gr.Button("Stop Training", variant="stop")
+                            training_status = gr.Textbox(label="Training Status", interactive=False)
+
+                        with gr.Column():
+                            progress_plot = gr.Plot(label="Training Progress")
+                            refresh_plot_button = gr.Button("Refresh Plot")
+
+                with gr.TabItem("4. Evaluation & Export"):
+                    with gr.Row():
+                        with gr.Column():
+                            test_prompt = gr.Textbox(
+                                label="Test Prompt",
+                                placeholder="Enter a prompt to test the model...",
+                                lines=3
+                            )
+                            max_gen_length = gr.Slider(
+                                minimum=10,
+                                maximum=500,
+                                step=10,
+                                value=100,
+                                label="Max Generation Length"
+                            )
+                            generate_button = gr.Button("Generate Text")
+                            generated_output = gr.Textbox(label="Generated Output", lines=10, interactive=False)
+
+                        with gr.Column():
+                            export_format = gr.Radio(
+                                ["pytorch", "tensorflow", "gguf"],
+                                label="Export Format",
+                                value="pytorch"
+                            )
+                            export_button = gr.Button("Export Model")
+                            export_status = gr.Textbox(label="Export Status", interactive=False)
+
+            # Functionality
+            def preprocess_data(file, format_type):
+                try:
+                    if file is None:
+                        return "Please upload a file first."
+
+                    # Process the uploaded file
+                    dataset = self.model_instance.prepare_dataset(file.name, format_type)
+                    self.model_instance.dataset = dataset
+
+                    # Create a summary of the dataset
+                    num_samples = len(dataset["train"])
+
+
+                    # Sample a few examples
+                    examples = dataset["train"].select(range(min(3, num_samples)))
+                    sample_text = []
+                    for ex in examples:
+                        text_key = list(ex.keys())[0] if "text" not in ex else "text"
+                        sample = ex[text_key]
+                        if isinstance(sample, str):
+                            sample_text.append(sample[:100] + "..." if len(sample) > 100 else sample)
+
+                    info = f"Dataset loaded successfully!\n"
+                    info += f"Number of training examples: {num_samples}\n"
+                    info += f"Sample data:\n" + "\n---\n".join(sample_text)
+
+                    return info
+                except Exception as e:
+                    return f"Error preprocessing data: {str(e)}"
+
+            def start_training(
+                model_name, learning_rate, batch_size, epochs, max_length,
+                use_lora, lora_r, lora_alpha, eval_ratio
+            ):
+                try:
+                    if self.model_instance.dataset is None:
+                        return "Please preprocess a dataset first."
+
+                    # Validate parameters
+                    if not model_name:
+                        return "Please select a model."
+                    
+                    # Prepare training parameters with proper type conversion
+                    training_params = {
+                        "model_name": str(model_name),
+                        "learning_rate": float(learning_rate),
+                        "batch_size": int(batch_size),
+                        "epochs": int(epochs),
+                        "max_length": int(max_length),
+                        "use_lora": bool(use_lora),
+                        "lora_r": int(lora_r) if use_lora else None,
+                        "lora_alpha": int(lora_alpha) if use_lora else None,
+                        "eval_ratio": float(eval_ratio),
+                        "weight_decay": float(self.default_params["weight_decay"]),
+                        "warmup_ratio": float(self.default_params["warmup_ratio"]),
+                        "lora_dropout": float(self.default_params["lora_dropout"])
+                    }
+
+                    # Start training in a separate thread
+                    import threading
+                    def train_thread():
+                        status = self.model_instance.train(training_params)
+                        return status
+
+                    thread = threading.Thread(target=train_thread)
+                    thread.start()
+
+                    return "Training started! Monitor the progress in the Training tab."
+                except Exception as e:
+                    return f"Error starting training: {str(e)}"
+
+            def stop_training():
+                if self.model_instance.trainer is not None:
+                    # Attempt to stop the trainer
+                    self.model_instance.trainer.stop_training = True
+                    return "Training stop signal sent. It may take a moment to complete the current step."
+                return "No active training to stop."
+
+            def update_progress_plot():
+                try:
+                    return self.model_instance.plot_training_progress()
+                except Exception as e:
+                    return None
+
+            def run_text_generation(prompt, max_length):
+                try:
+                    if self.model_instance.model is None:
+                        return "Please fine-tune a model first."
+
+                    return self.model_instance.generate_text(prompt, int(max_length))
+                except Exception as e:
+                    return f"Error generating text: {str(e)}"
+
+            def export_model_fn(format_type):
+                try:
+                    if self.model_instance.model is None:
+                        return "Please fine-tune a model first."
+
+                    return self.model_instance.export_model(format_type)
+                except Exception as e:
+                    return f"Error exporting model: {str(e)}"
+
+            # Connect UI components to functions
+            preprocess_button.click(
+                preprocess_data,
+                inputs=[file_upload, file_format],
+                outputs=dataset_info
+            )
+
+            start_training_button.click(
+                start_training,
+                inputs=[
+                    model_name, learning_rate, batch_size, epochs, max_length,
+                    use_lora, lora_r, lora_alpha, eval_ratio
+                ],
+                outputs=training_status
+            )
+
+            stop_training_button.click(
+                stop_training,
+                inputs=[],
+                outputs=training_status
+            )
+
+            refresh_plot_button.click(
+                update_progress_plot,
+                inputs=[],
+                outputs=progress_plot
+            )
+
+            generate_button.click(
+                run_text_generation,
+                inputs=[test_prompt, max_gen_length],
+                outputs=generated_output
+            )
+
+            export_button.click(
+                export_model_fn,
+                inputs=[export_format],
+                outputs=export_status
+            )
+
+        return app
+
+if __name__ == '__main__':
+    ui = GemmaUI()
+    app = ui.create_ui()
+    app.launch()