Create app.py

app.py

@@ -0,0 +1,160 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from datasets import load_dataset
+import torch
+import os
+
+#-------------------------------Functions----------------------------------------------#
+def load_and_preprocess_data(dataset_name, tokenizer):
+    try:
+        dataset = load_dataset(dataset_name, split="train")
+    except Exception as e:
+        return None, f"Error loading dataset: {e}"
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
+
+    try:
+        tokenized_datasets = dataset.map(tokenize_function, batched=True, num_proc=4, remove_columns=["text"])
+    except Exception as e:
+        return None, f"Error tokenizing dataset: {e}"
+
+    return tokenized_datasets, None
+#---------------------------------------------------------------------------------------#
+
+def train_model(architecture_size, api_key, repo_name, push_to_hub):
+    # Map architecture size to model name
+    model_name_mapping = {
+        "Small": "distilgpt2",
+        "Medium": "gpt2",
+        "Large": "gpt2-medium",
+    }
+    model_name = model_name_mapping[architecture_size]
+
+    # Device setup
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device_msg = "CUDA is available! Training will be faster on GPU." if torch.cuda.is_available() else "CUDA not available. Training on CPU will be slow."
+
+    # Validate push_to_hub inputs
+    if push_to_hub:
+        if not api_key or not api_key.strip():
+            return "❌ Error: You must provide a Hugging Face API key if pushing to hub is selected."
+        if not repo_name or not repo_name.strip():
+            return "❌ Error: You must provide a repository name if pushing to hub is selected."
+
+    try:
+        # Load dataset
+        dataset_name = "wikitext-2-raw-v1"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        tokenized_datasets, error_msg = load_and_preprocess_data(dataset_name, tokenizer)
+        if error_msg:
+            return f"❌ {error_msg}"
+        if tokenized_datasets is None:
+            return "❌ Failed to load and preprocess dataset."
+
+        # Load model
+        model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+        model.resize_token_embeddings(len(tokenizer))
+
+        # Training args
+        output_dir = "./results"
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=1,
+            per_device_train_batch_size=4,
+            save_steps=500,
+            save_total_limit=1,
+            logging_steps=250,
+            learning_rate=5e-5,
+            weight_decay=0.01,
+            push_to_hub=push_to_hub,
+            hub_model_id=repo_name if push_to_hub else None,
+            hub_token=api_key if push_to_hub else None,
+            fp16=torch.cuda.is_available(),
+        )
+
+        # Data collator
+        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+        # Trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_datasets,
+            data_collator=data_collator,
+        )
+
+        # Train
+        trainer.train()
+
+        # Save locally
+        trainer.save_model(output_dir)
+
+        # Evaluate
+        eval_results = trainer.evaluate()
+        eval_loss = eval_results.get('eval_loss', 'N/A')
+
+        # Push to hub if selected
+        if push_to_hub:
+            trainer.push_to_hub()
+            hub_msg = f"✅ Model pushed to Hugging Face Hub: {repo_name}"
+        else:
+            hub_msg = "ℹ️ Model saved locally at ./results (not pushed to hub)."
+
+        return f"""✅ Training Complete!
+- Device: {device_msg}
+- Eval Loss: {eval_loss}
+- {hub_msg}
+"""
+
+    except Exception as e:
+        return f"❌ Training Error: {str(e)}"
+
+# ----------------------------- Gradio Interface ----------------------------- #
+
+with gr.Blocks(title="LLM Builder - Gradio") as demo:
+    gr.Markdown("# 🤖 LLM Builder")
+    gr.Markdown("### 1. Select Model Architecture")
+
+    architecture_size = gr.Dropdown(
+        choices=["Small", "Medium", "Large"],
+        value="Small",
+        label="Choose Model Size",
+        info="Select the size of the model. Larger models have more parameters."
+    )
+
+    gr.Markdown("### 2. Training Setup")
+
+    with gr.Row():
+        with gr.Column():
+            api_key = gr.Textbox(
+                label="Hugging Face Hub API Key",
+                type="password",
+                placeholder="hf_...",
+                info="Required only if pushing to hub."
+            )
+            repo_name = gr.Textbox(
+                label="Repository Name",
+                placeholder="your-username/your-model-name",
+                info="Required only if pushing to hub."
+            )
+            push_to_hub = gr.Checkbox(
+                label="Push to Hugging Face Hub?",
+                value=False
+            )
+
+    train_btn = gr.Button("🚀 Start Training", variant="primary")
+    output = gr.Textbox(label="Training Output", placeholder="Training logs and results will appear here...", lines=10)
+
+    train_btn.click(
+        fn=train_model,
+        inputs=[architecture_size, api_key, repo_name, push_to_hub],
+        outputs=output
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()