{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# PromptWizard Qwen Fine-tuning on HuggingFace\n", "\n", "This notebook fine-tunes Qwen models using GSM8K dataset with PromptWizard methodology.\n", "Run this on HuggingFace or Google Colab with GPU support." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Install required packages\n", "!pip install -q transformers==4.36.2 datasets==2.16.1 peft==0.7.1 accelerate==0.25.0 bitsandbytes==0.41.3" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import json\n", "from datasets import Dataset, load_dataset\n", "from transformers import (\n", " AutoModelForCausalLM,\n", " AutoTokenizer,\n", " TrainingArguments,\n", " Trainer,\n", " DataCollatorForLanguageModeling\n", ")\n", "from peft import LoraConfig, get_peft_model, TaskType\n", "\n", "# Check GPU availability\n", "if torch.cuda.is_available():\n", " print(f\"✅ GPU Available: {torch.cuda.get_device_name(0)}\")\n", " print(f\" Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")\n", "else:\n", " print(\"⚠️ No GPU detected. Training will be slow.\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Load and prepare GSM8K dataset\n", "print(\"Loading GSM8K dataset...\")\n", "dataset = load_dataset(\"openai/gsm8k\", \"main\")\n", "\n", "def format_prompt(item):\n", " \"\"\"Format GSM8K item for training\"\"\"\n", " prompt = f\"\"\"<|system|>\n", "You are a mathematics expert. Solve grade school math problems step by step.\n", "<|user|>\n", "{item['question']}\n", "<|assistant|>\n", "Let me solve this step by step.\n", "\n", "{item['answer']}\"\"\"\n", " return {\"text\": prompt}\n", "\n", "# Process datasets (using smaller subset for demo)\n", "train_data = dataset['train'].select(range(min(1000, len(dataset['train']))))\n", "eval_data = dataset['test'].select(range(min(100, len(dataset['test']))))\n", "\n", "train_dataset = train_data.map(format_prompt)\n", "eval_dataset = eval_data.map(format_prompt)\n", "\n", "print(f\"Train samples: {len(train_dataset)}\")\n", "print(f\"Eval samples: {len(eval_dataset)}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Load model and tokenizer\n", "MODEL_NAME = \"Qwen/Qwen2.5-1.5B\" # Using smaller model for faster training\n", "\n", "print(f\"Loading {MODEL_NAME}...\")\n", "\n", "tokenizer = AutoTokenizer.from_pretrained(\n", " MODEL_NAME,\n", " trust_remote_code=True,\n", " padding_side=\"left\"\n", ")\n", "\n", "if tokenizer.pad_token is None:\n", " tokenizer.pad_token = tokenizer.eos_token\n", "\n", "model = AutoModelForCausalLM.from_pretrained(\n", " MODEL_NAME,\n", " trust_remote_code=True,\n", " torch_dtype=torch.float16,\n", " device_map=\"auto\",\n", " load_in_8bit=True\n", ")\n", "\n", "# Configure LoRA\n", "lora_config = LoraConfig(\n", " task_type=TaskType.CAUSAL_LM,\n", " r=8, # Lower rank for faster training\n", " lora_alpha=16,\n", " lora_dropout=0.1,\n", " target_modules=[\"q_proj\", \"v_proj\"],\n", " bias=\"none\"\n", ")\n", "\n", "model = get_peft_model(model, lora_config)\n", "model.print_trainable_parameters()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Tokenize datasets\n", "def tokenize_function(examples):\n", " return tokenizer(\n", " examples[\"text\"],\n", " padding=\"max_length\",\n", " truncation=True,\n", " max_length=512\n", " )\n", "\n", "print(\"Tokenizing datasets...\")\n", "train_dataset = train_dataset.map(tokenize_function, batched=True)\n", "eval_dataset = eval_dataset.map(tokenize_function, batched=True)\n", "\n", "# Data collator\n", "data_collator = DataCollatorForLanguageModeling(\n", " tokenizer=tokenizer,\n", " mlm=False,\n", " pad_to_multiple_of=8\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Training arguments\n", "training_args = TrainingArguments(\n", " output_dir=\"./qwen-promptwizard\",\n", " num_train_epochs=1, # Quick training for demo\n", " per_device_train_batch_size=4,\n", " per_device_eval_batch_size=4,\n", " gradient_accumulation_steps=4,\n", " warmup_steps=100,\n", " logging_steps=10,\n", " save_steps=100,\n", " evaluation_strategy=\"steps\",\n", " eval_steps=50,\n", " save_total_limit=2,\n", " load_best_model_at_end=True,\n", " fp16=True,\n", " push_to_hub=False, # Set to True to push to HF Hub\n", " gradient_checkpointing=True,\n", " optim=\"adamw_torch\",\n", " learning_rate=2e-4,\n", " lr_scheduler_type=\"cosine\",\n", ")\n", "\n", "# Initialize trainer\n", "trainer = Trainer(\n", " model=model,\n", " args=training_args,\n", " train_dataset=train_dataset,\n", " eval_dataset=eval_dataset,\n", " data_collator=data_collator,\n", " tokenizer=tokenizer,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Start training\n", "print(\"Starting training...\")\n", "print(f\"Using {torch.cuda.device_count()} GPU(s)\")\n", "\n", "trainer.train()\n", "\n", "print(\"\\n✅ Training complete!\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Save model\n", "print(\"Saving model...\")\n", "trainer.save_model(\"./qwen-promptwizard-final\")\n", "tokenizer.save_pretrained(\"./qwen-promptwizard-final\")\n", "\n", "print(\"Model saved to ./qwen-promptwizard-final\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Test the fine-tuned model\n", "from transformers import pipeline\n", "\n", "# Load the fine-tuned model\n", "generator = pipeline(\n", " \"text-generation\",\n", " model=\"./qwen-promptwizard-final\",\n", " tokenizer=tokenizer,\n", " device_map=\"auto\"\n", ")\n", "\n", "# Test prompt\n", "test_prompt = \"\"\"<|system|>\n", "You are a mathematics expert. Solve grade school math problems step by step.\n", "<|user|>\n", "Janet has 3 apples. She buys 5 more apples from the store. How many apples does she have now?\n", "<|assistant|>\"\"\"\n", "\n", "# Generate response\n", "response = generator(\n", " test_prompt,\n", " max_new_tokens=200,\n", " temperature=0.7,\n", " do_sample=True\n", ")\n", "\n", "print(\"Test Response:\")\n", "print(response[0]['generated_text'])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Next Steps\n", "\n", "1. **Push to HuggingFace Hub**: Set `push_to_hub=True` in training arguments\n", "2. **Increase Training**: Use more epochs and larger dataset for better results\n", "3. **Use Larger Model**: Try Qwen2.5-7B for better performance (needs more GPU memory)\n", "4. **Fine-tune Hyperparameters**: Adjust learning rate, LoRA rank, etc.\n", "\n", "The trained model can now be used with PromptWizard for enhanced prompt optimization!" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 4 }