Spaces:
Sleeping
Sleeping
| """ | |
| PromptWizard — Qwen2.5-0.5B Fine-tuning on Bhagavad Gita Dataset | |
| Downloads CSV locally before training (for Hugging Face Spaces) | |
| """ | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| import pandas as pd | |
| from datasets import Dataset | |
| from transformers import ( | |
| AutoTokenizer, | |
| AutoModelForCausalLM, | |
| Trainer, | |
| TrainingArguments, | |
| ) | |
| from peft import LoraConfig, get_peft_model, TaskType | |
| from huggingface_hub import snapshot_download, HfApi | |
| import os | |
| import shutil | |
| # ------------------------------------------------------ | |
| # 🧠 GPU check | |
| # ------------------------------------------------------ | |
| def check_gpu_status(): | |
| return "🟢 Ready — GPU will be assigned at runtime (Zero GPU mode)" | |
| # ------------------------------------------------------ | |
| # 🧩 Download Dataset to /tmp/ | |
| # ------------------------------------------------------ | |
| def download_gita_dataset(): | |
| repo_id = "rahul7star/Gita" | |
| local_dir = "/tmp/gita_data" | |
| if os.path.exists(local_dir): | |
| shutil.rmtree(local_dir) | |
| os.makedirs(local_dir, exist_ok=True) | |
| print(f"📥 Downloading dataset from {repo_id} ...") | |
| snapshot_download(repo_id=repo_id, local_dir=local_dir, repo_type="dataset") | |
| # Try to locate the CSV file | |
| csv_path = None | |
| for root, _, files in os.walk(local_dir): | |
| for f in files: | |
| if f.lower().endswith(".csv"): | |
| csv_path = os.path.join(root, f) | |
| break | |
| if not csv_path: | |
| raise FileNotFoundError("No CSV file found in the Gita dataset repository.") | |
| print(f"✅ Found CSV: {csv_path}") | |
| return csv_path | |
| # ------------------------------------------------------ | |
| # 🚀 Training function | |
| # ------------------------------------------------------ | |
| def train_model(model_name, num_epochs, batch_size, learning_rate, progress=gr.Progress()): | |
| logs = [] | |
| try: | |
| progress(0.05, desc="Initializing...") | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| logs.append(f"🎮 Device: {device}") | |
| # ------------------------------------------------------ | |
| # 📂 Step 1: Download dataset | |
| # ------------------------------------------------------ | |
| progress(0.1, desc="Downloading dataset...") | |
| logs.append("\n📥 Downloading Gita dataset from HF Hub...") | |
| csv_path = download_gita_dataset() | |
| # ------------------------------------------------------ | |
| # 🧾 Step 2: Load dataset from CSV | |
| # ------------------------------------------------------ | |
| progress(0.2, desc="Loading dataset...") | |
| df = pd.read_csv(csv_path) | |
| if not all(c in df.columns for c in ["question", "answer"]): | |
| raise ValueError("CSV must contain 'question' and 'answer' columns.") | |
| hf_dataset = Dataset.from_pandas(df) | |
| def format_row(row): | |
| return { | |
| "text": f"<|system|>\nYou are a spiritual guide explaining Gita concepts.\n" | |
| f"<|user|>\n{row['question']}\n" | |
| f"<|assistant|>\n{row['answer']}" | |
| } | |
| hf_dataset = hf_dataset.map(format_row) | |
| logs.append(f"✅ Loaded {len(hf_dataset)} examples from {csv_path}") | |
| # ------------------------------------------------------ | |
| # 🤖 Step 3: Load model + tokenizer | |
| # ------------------------------------------------------ | |
| progress(0.3, desc="Loading Qwen model...") | |
| model_name = "Qwen/Qwen2.5-0.5B" | |
| logs.append(f"\n🔍 Loading base model: {model_name}") | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| torch_dtype=torch.float16 if device == "cuda" else torch.float32, | |
| ) | |
| if device == "cuda": | |
| model = model.to(device) | |
| logs.append("✅ Model and tokenizer ready") | |
| # ------------------------------------------------------ | |
| # ⚙️ Step 4: Apply LoRA config | |
| # ------------------------------------------------------ | |
| progress(0.4, desc="Configuring LoRA...") | |
| lora_config = LoraConfig( | |
| task_type=TaskType.CAUSAL_LM, | |
| r=8, | |
| lora_alpha=16, | |
| lora_dropout=0.1, | |
| target_modules=["q_proj", "v_proj"], | |
| ) | |
| model = get_peft_model(model, lora_config) | |
| # ------------------------------------------------------ | |
| # ✂️ Step 5: Tokenize dataset | |
| # ------------------------------------------------------ | |
| progress(0.5, desc="Tokenizing data...") | |
| def tokenize_fn(batch): | |
| return tokenizer( | |
| batch["text"], | |
| truncation=True, | |
| padding="max_length", | |
| max_length=512, | |
| ) | |
| tokenized = hf_dataset.map(tokenize_fn, batched=True) | |
| logs.append("🧾 Dataset tokenized successfully") | |
| # ------------------------------------------------------ | |
| # 🎯 Step 6: Training arguments | |
| # ------------------------------------------------------ | |
| progress(0.6, desc="Configuring training...") | |
| training_args = TrainingArguments( | |
| output_dir="/tmp/qwen-gita-output", | |
| num_train_epochs=num_epochs, | |
| per_device_train_batch_size=batch_size, | |
| gradient_accumulation_steps=2, | |
| learning_rate=learning_rate, | |
| warmup_steps=5, | |
| logging_steps=5, | |
| save_strategy="no", | |
| fp16=device == "cuda", | |
| max_steps=100, | |
| report_to="none", | |
| ) | |
| # ------------------------------------------------------ | |
| # 🏋️ Step 7: Train model | |
| # ------------------------------------------------------ | |
| progress(0.7, desc="Training in progress...") | |
| logs.append("\n🚀 Starting fine-tuning...") | |
| trainer = Trainer( | |
| model=model, | |
| args=training_args, | |
| train_dataset=tokenized, | |
| tokenizer=tokenizer, | |
| ) | |
| trainer.train() | |
| # ------------------------------------------------------ | |
| # 💾 Step 8: Save + Upload | |
| # ------------------------------------------------------ | |
| progress(0.9, desc="Saving and uploading...") | |
| output_dir = "/tmp/qwen-gita-lora" | |
| os.makedirs(output_dir, exist_ok=True) | |
| model.save_pretrained(output_dir) | |
| tokenizer.save_pretrained(output_dir) | |
| logs.append("\n📤 Uploading fine-tuned LoRA model to Hugging Face Hub...") | |
| repo_id = "rahul7star/Qwen0.5-3B-Gita" | |
| api = HfApi() | |
| api.upload_folder(folder_path=output_dir, repo_id=repo_id) | |
| logs.append(f"✅ Uploaded fine-tuned model to {repo_id}") | |
| progress(1.0, desc="Complete!") | |
| logs.append("\n🎉 Training complete!") | |
| except Exception as e: | |
| logs.append(f"\n❌ Error: {str(e)}") | |
| return "\n".join(logs) | |
| # ------------------------------------------------------ | |
| # 🎨 Gradio Interface | |
| # ------------------------------------------------------ | |
| def create_interface(): | |
| with gr.Blocks(title="PromptWizard — Qwen Gita Trainer") as demo: | |
| gr.Markdown(""" | |
| # 🧘 PromptWizard — Qwen2.5-0.5B Gita Trainer | |
| Fine-tune Qwen 0.5B on your **Bhagavad Gita CSV dataset** | |
| Automatically uploads LoRA weights to `rahul7star/Qwen0.5-3B-Gita` | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gpu_status = gr.Textbox( | |
| label="GPU Status", | |
| value=check_gpu_status(), | |
| interactive=False, | |
| ) | |
| model_name = gr.Textbox( | |
| value="Qwen/Qwen2.5-0.5B", | |
| visible=False, | |
| ) | |
| num_epochs = gr.Slider(1, 3, 1, step=1, label="Epochs") | |
| batch_size = gr.Slider(1, 4, 2, step=1, label="Batch Size") | |
| learning_rate = gr.Number(value=5e-5, label="Learning Rate") | |
| train_btn = gr.Button("🚀 Start Fine-tuning", variant="primary") | |
| with gr.Column(): | |
| output = gr.Textbox( | |
| label="Training Logs", | |
| lines=25, | |
| max_lines=40, | |
| value="Click 'Start Fine-tuning' to train on Bhagavad Gita dataset...", | |
| ) | |
| train_btn.click( | |
| fn=train_model, | |
| inputs=[model_name, num_epochs, batch_size, learning_rate], | |
| outputs=output, | |
| ) | |
| gr.Markdown(""" | |
| --- | |
| **Notes:** | |
| - Downloads dataset: `rahul7star/Gita` → `/tmp/gita_data/Gita.csv` | |
| - Trains using LoRA for efficiency | |
| - Uploads to `rahul7star/Qwen0.5-3B-Gita` | |
| """) | |
| return demo | |
| # ------------------------------------------------------ | |
| # 🚪 Launch app | |
| # ------------------------------------------------------ | |
| if __name__ == "__main__": | |
| demo = create_interface() | |
| demo.launch() |