PROPER MISTRAL 7B TRAINING: v0.3 with LoRA, proper chat format, no logging issues

CELESTIAL_Training_Notebook.ipynb

@@ -4,14 +4,14 @@
             "cell_type": "markdown",
             "metadata": {},
             "source": [
-                "# 🌟 CELESTIAL PRODUCTION TRAINING\n",
-                "## 150 Perfect Conversations - Production Ready\n",
-                "\n",
-                "This notebook trains CELESTIAL AI with production-quality conversations:\n",
-                "- 100 comprehensive numerology conversations\n",
-                "- 50 authentic Krishna divine guidance conversations\n",
-                "- Each response is perfect, coherent, and detailed\n",
-                "- No wandb issues - clean, reliable training"
+                "# 🌟 CELESTIAL MISTRAL 7B TRAINING\n",
+                "## Train Your Own Mistral 7B Model for CELESTIAL AI\n",
+                "\n",
+                "This notebook properly trains Mistral 7B v0.3 with:\n",
+                "- 150 production-quality conversations\n",
+                "- LoRA fine-tuning for efficiency\n",
+                "- Proper chat formatting for Mistral\n",
+                "- No logging issues"
             ]
         },
         {
@@ -20,17 +20,19 @@
             "metadata": {},
             "outputs": [],
             "source": [
-                "# 📦 INSTALL AND SETUP\n",
-                "!pip install -q transformers datasets accelerate peft bitsandbytes huggingface_hub\n",
+                "# 📦 INSTALL REQUIRED PACKAGES FOR MISTRAL 7B\n",
+                "!pip install -q transformers==4.36.0 datasets accelerate peft bitsandbytes huggingface_hub trl\n",
                 "\n",
-                "# Disable all logging that might cause issues\n",
+                "# Disable all logging to prevent issues\n",
                 "import os\n",
                 "import warnings\n",
                 "os.environ[\"WANDB_DISABLED\"] = \"true\"\n",
                 "os.environ[\"WANDB_MODE\"] = \"disabled\"\n",
+                "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
                 "warnings.filterwarnings('ignore')\n",
                 "\n",
-                "print('✅ All packages installed and logging disabled!')"
+                "print('✅ Packages installed for Mistral 7B training!')\n",
+                "print('🚫 All logging disabled to prevent errors')"
             ]
         },
         {
@@ -42,13 +44,13 @@
                 "# 🔑 HUGGINGFACE AUTHENTICATION\n",
                 "from huggingface_hub import notebook_login\n",
                 "\n",
-                "print('🔐 Authenticating with HuggingFace...')\n",
+                "print('🔐 Authenticating with HuggingFace for Mistral access...')\n",
                 "try:\n",
                 "    notebook_login()\n",
                 "    print('✅ Authentication successful!')\n",
                 "except Exception as e:\n",
                 "    print(f'⚠️ Authentication failed: {e}')\n",
-                "    print('Please manually set your HF token if needed')"
+                "    print('Please set your HF token manually if needed')"
             ]
         },
         {
@@ -57,20 +59,19 @@
             "metadata": {},
             "outputs": [],
             "source": [
-                "# 📊 LOAD PRODUCTION DATASET\n",
+                "# 📊 LOAD CELESTIAL DATASET\n",
                 "from datasets import load_dataset\n",
                 "\n",
                 "DATASET_REPO = 'dp1812/celestial-comprehensive-spiritual-ai'\n",
                 "\n",
-                "print('📊 Loading PRODUCTION dataset...')\n",
+                "print('📊 Loading CELESTIAL dataset for Mistral training...')\n",
                 "try:\n",
                 "    dataset = load_dataset(DATASET_REPO, data_files='celestial_complete_production_dataset.jsonl', split='train')\n",
-                "    print(f'✅ Dataset loaded: {len(dataset)} production-quality conversations')\n",
+                "    print(f'✅ Dataset loaded: {len(dataset)} conversations')\n",
                 "    print('🎯 100 numerology + 50 Krishna divine guidance')\n",
-                "    print('💎 Each conversation is perfect and coherent!')\n",
                 "except Exception as e:\n",
                 "    print(f'❌ Dataset loading failed: {e}')\n",
-                "    # Fallback to main dataset\n",
+                "    # Fallback\n",
                 "    try:\n",
                 "        dataset = load_dataset(DATASET_REPO, split='train')\n",
                 "        print(f'✅ Fallback dataset loaded: {len(dataset)} conversations')\n",
@@ -81,8 +82,8 @@
                 "# Show sample\n",
                 "print('\\n📝 Sample conversation:')\n",
                 "sample = dataset[0]\n",
-                "print(f\"User: {sample['messages'][1]['content'][:80]}...\")\n",
-                "print(f\"Assistant: {sample['messages'][2]['content'][:80]}...\")"
+                "print(f\"User: {sample['messages'][1]['content'][:60]}...\")\n",
+                "print(f\"Assistant: {sample['messages'][2]['content'][:60]}...\")"
             ]
         },
         {
@@ -91,20 +92,27 @@
             "metadata": {},
             "outputs": [],
             "source": [
-                "# 🤖 LOAD MODEL AND TOKENIZER\n",
+                "# 🤖 LOAD MISTRAL 7B MODEL AND TOKENIZER\n",
                 "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
                 "import torch\n",
                 "\n",
-                "MODEL_NAME = 'microsoft/DialoGPT-medium'\n",
+                "MODEL_NAME = 'mistralai/Mistral-7B-v0.3'\n",
                 "\n",
-                "print('🤖 Loading model and tokenizer...')\n",
+                "print('🤖 Loading Mistral 7B v0.3 model and tokenizer...')\n",
                 "\n",
-                "# Load tokenizer\n",
-                "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)\n",
+                "# Load tokenizer with proper settings\n",
+                "tokenizer = AutoTokenizer.from_pretrained(\n",
+                "    MODEL_NAME,\n",
+                "    trust_remote_code=True,\n",
+                "    padding_side='right'\n",
+                ")\n",
+                "\n",
+                "# Add pad token if missing\n",
                 "if tokenizer.pad_token is None:\n",
                 "    tokenizer.pad_token = tokenizer.eos_token\n",
+                "    tokenizer.pad_token_id = tokenizer.eos_token_id\n",
                 "\n",
-                "# Load model with quantization for efficiency\n",
+                "# Quantization config for efficient training\n",
                 "bnb_config = BitsAndBytesConfig(\n",
                 "    load_in_4bit=True,\n",
                 "    bnb_4bit_quant_type=\"nf4\",\n",
@@ -112,16 +120,19 @@
                 "    bnb_4bit_use_double_quant=True\n",
                 ")\n",
                 "\n",
+                "# Load Mistral 7B model\n",
                 "model = AutoModelForCausalLM.from_pretrained(\n",
                 "    MODEL_NAME,\n",
                 "    quantization_config=bnb_config,\n",
                 "    device_map=\"auto\",\n",
-                "    trust_remote_code=True\n",
+                "    trust_remote_code=True,\n",
+                "    torch_dtype=torch.float16\n",
                 ")\n",
                 "\n",
-                "print('✅ Model and tokenizer loaded successfully!')\n",
-                "print(f'🔍 Model type: {type(model)}')\n",
-                "print(f'🔍 Tokenizer vocab size: {len(tokenizer)}')"
+                "print('✅ Mistral 7B model and tokenizer loaded successfully!')\n",
+                "print(f'🔍 Model: {MODEL_NAME}')\n",
+                "print(f'🔍 Tokenizer vocab size: {len(tokenizer)}')\n",
+                "print(f'🔍 Model device: {model.device}')"
             ]
         },
         {
@@ -130,44 +141,48 @@
             "metadata": {},
             "outputs": [],
             "source": [
-                "# 🔧 SETUP LORA FOR EFFICIENT TRAINING\n",
-                "from peft import LoraConfig, get_peft_model, TaskType\n",
+                "# 🔧 SETUP LORA FOR MISTRAL 7B\n",
+                "from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training\n",
                 "\n",
-                "print('🔧 Setting up LoRA for efficient training...')\n",
+                "print('🔧 Setting up LoRA for Mistral 7B training...')\n",
                 "\n",
-                "# Auto-detect target modules\n",
-                "def find_target_modules(model):\n",
-                "    target_modules = set()\n",
-                "    for name, module in model.named_modules():\n",
-                "        if isinstance(module, torch.nn.Linear):\n",
-                "            module_name = name.split('.')[-1]\n",
-                "            if any(pattern in module_name for pattern in ['attn', 'proj', 'fc', 'dense']):\n",
-                "                target_modules.add(module_name)\n",
-                "    return list(target_modules) if target_modules else ['c_attn', 'c_proj']\n",
+                "# Prepare model for k-bit training\n",
+                "model = prepare_model_for_kbit_training(model)\n",
                 "\n",
-                "target_modules = find_target_modules(model)\n",
-                "print(f'🎯 Target modules detected: {target_modules}')\n",
+                "# Mistral 7B specific target modules\n",
+                "target_modules = [\n",
+                "    \"q_proj\",\n",
+                "    \"k_proj\", \n",
+                "    \"v_proj\",\n",
+                "    \"o_proj\",\n",
+                "    \"gate_proj\",\n",
+                "    \"up_proj\",\n",
+                "    \"down_proj\",\n",
+                "    \"lm_head\"\n",
+                "]\n",
                 "\n",
-                "# Create LoRA config\n",
+                "print(f'🎯 Target modules for Mistral: {target_modules}')\n",
+                "\n",
+                "# Create LoRA config optimized for Mistral\n",
                 "lora_config = LoraConfig(\n",
-                "    r=16,\n",
-                "    lora_alpha=32,\n",
+                "    r=64,  # Higher rank for better performance\n",
+                "    lora_alpha=16,\n",
                 "    target_modules=target_modules,\n",
                 "    lora_dropout=0.1,\n",
                 "    bias=\"none\",\n",
                 "    task_type=TaskType.CAUSAL_LM,\n",
                 ")\n",
                 "\n",
-                "# Apply LoRA with error handling\n",
+                "# Apply LoRA to Mistral\n",
                 "try:\n",
                 "    model = get_peft_model(model, lora_config)\n",
                 "    model.print_trainable_parameters()\n",
-                "    print('✅ LoRA adapters attached successfully!')\n",
+                "    print('✅ LoRA adapters attached to Mistral 7B!')\n",
                 "except Exception as e:\n",
-                "    print(f'⚠️ LoRA failed: {e}')\n",
-                "    print('🔧 Continuing with full fine-tuning')\n",
+                "    print(f'❌ LoRA setup failed: {e}')\n",
+                "    raise\n",
                 "\n",
-                "print('🎯 Model ready for production training!')"
+                "print('🎯 Mistral 7B ready for CELESTIAL training!')"
             ]
         },
         {
@@ -176,40 +191,42 @@
             "metadata": {},
             "outputs": [],
             "source": [
-                "# 📝 PREPARE TRAINING DATA\n",
-                "def format_conversation(example):\n",
-                "    \"\"\"Format conversation for training\"\"\"\n",
+                "# 📝 FORMAT DATA FOR MISTRAL CHAT TRAINING\n",
+                "def format_for_mistral_chat(example):\n",
+                "    \"\"\"Format conversation for Mistral chat training\"\"\"\n",
                 "    messages = example['messages']\n",
                 "    \n",
-                "    # Extract user and assistant messages\n",
+                "    # Extract messages\n",
+                "    system_msg = messages[0]['content']\n",
                 "    user_msg = messages[1]['content']\n",
                 "    assistant_msg = messages[2]['content']\n",
                 "    \n",
-                "    # Create training format\n",
-                "    formatted = f\"User: {user_msg}\\nCELESTIAL AI: {assistant_msg}<|endoftext|>\"\n",
+                "    # Mistral chat format\n",
+                "    formatted = f\"<s>[INST] {system_msg}\\n\\nUser: {user_msg} [/INST] {assistant_msg}</s>\"\n",
                 "    \n",
-                "    # Tokenize with proper settings\n",
+                "    # Tokenize\n",
                 "    tokens = tokenizer(\n",
                 "        formatted,\n",
                 "        truncation=True,\n",
-                "        padding='max_length',\n",
-                "        max_length=1024,  # Longer for detailed responses\n",
-                "        return_tensors='pt'\n",
+                "        padding=False,\n",
+                "        max_length=2048,  # Mistral context length\n",
+                "        return_tensors=None\n",
                 "    )\n",
                 "    \n",
-                "    # Set labels for training\n",
-                "    tokens['labels'] = tokens['input_ids'].clone()\n",
+                "    # Set labels (same as input_ids for causal LM)\n",
+                "    tokens['labels'] = tokens['input_ids'].copy()\n",
                 "    \n",
-                "    return {\n",
-                "        'input_ids': tokens['input_ids'].squeeze(),\n",
-                "        'attention_mask': tokens['attention_mask'].squeeze(),\n",
-                "        'labels': tokens['labels'].squeeze()\n",
-                "    }\n",
-                "\n",
-                "print('📝 Formatting production training data...')\n",
-                "formatted_dataset = dataset.map(format_conversation, remove_columns=dataset.column_names)\n",
-                "print(f'✅ Formatted {len(formatted_dataset)} conversations for training')\n",
-                "print('🎯 Each conversation is optimized for CELESTIAL AI responses')"
+                "    return tokens\n",
+                "\n",
+                "print('📝 Formatting data for Mistral chat training...')\n",
+                "formatted_dataset = dataset.map(\n",
+                "    format_for_mistral_chat,\n",
+                "    remove_columns=dataset.column_names,\n",
+                "    desc=\"Formatting for Mistral\"\n",
+                ")\n",
+                "\n",
+                "print(f'✅ Formatted {len(formatted_dataset)} conversations for Mistral')\n",
+                "print('🎯 Using proper Mistral chat format with [INST] tags')"
             ]
         },
         {
@@ -218,38 +235,40 @@
             "metadata": {},
             "outputs": [],
             "source": [
-                "# 🚀 PRODUCTION TRAINING CONFIGURATION\n",
+                "# 🚀 MISTRAL TRAINING CONFIGURATION\n",
                 "from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling\n",
                 "\n",
-                "print('🚀 Setting up production training configuration...')\n",
+                "print('🚀 Setting up Mistral 7B training configuration...')\n",
                 "\n",
-                "# Training arguments optimized for production\n",
+                "# Training arguments optimized for Mistral 7B\n",
                 "training_args = TrainingArguments(\n",
-                "    output_dir='./celestial-production-results',\n",
-                "    num_train_epochs=5,  # More epochs for better learning\n",
-                "    per_device_train_batch_size=1,  # Conservative for stability\n",
-                "    gradient_accumulation_steps=8,  # Effective batch size of 8\n",
-                "    warmup_steps=20,\n",
-                "    learning_rate=3e-5,  # Conservative learning rate\n",
+                "    output_dir='./celestial-mistral-7b-results',\n",
+                "    num_train_epochs=3,\n",
+                "    per_device_train_batch_size=1,\n",
+                "    gradient_accumulation_steps=16,  # Effective batch size of 16\n",
+                "    warmup_steps=50,\n",
+                "    learning_rate=2e-4,  # Higher LR for LoRA\n",
                 "    fp16=True,\n",
                 "    logging_steps=10,\n",
-                "    save_steps=50,\n",
-                "    evaluation_strategy='no',\n",
+                "    save_steps=100,\n",
+                "    eval_strategy='no',\n",
                 "    save_strategy='steps',\n",
                 "    load_best_model_at_end=False,\n",
                 "    report_to=[],  # No external logging\n",
                 "    remove_unused_columns=False,\n",
                 "    dataloader_drop_last=True,\n",
-                "    disable_tqdm=False\n",
+                "    group_by_length=True,  # Efficient batching\n",
+                "    ddp_find_unused_parameters=False\n",
                 ")\n",
                 "\n",
-                "# Data collator for language modeling\n",
+                "# Data collator for Mistral\n",
                 "data_collator = DataCollatorForLanguageModeling(\n",
                 "    tokenizer=tokenizer,\n",
-                "    mlm=False  # Causal LM, not masked LM\n",
+                "    mlm=False,\n",
+                "    pad_to_multiple_of=8\n",
                 ")\n",
                 "\n",
-                "# Create trainer\n",
+                "# Create Mistral trainer\n",
                 "trainer = Trainer(\n",
                 "    model=model,\n",
                 "    args=training_args,\n",
@@ -258,9 +277,9 @@
                 "    data_collator=data_collator\n",
                 ")\n",
                 "\n",
-                "print('✅ Production training configuration ready!')\n",
-                "print('🎯 Optimized for high-quality CELESTIAL AI training')\n",
-                "print('⏱️ Expected training time: 20-30 minutes')"
+                "print('✅ Mistral 7B training configuration ready!')\n",
+                "print('🎯 Optimized for CELESTIAL AI with LoRA fine-tuning')\n",
+                "print('⏱️ Expected training time: 30-45 minutes')"
             ]
         },
         {
@@ -269,23 +288,23 @@
             "metadata": {},
             "outputs": [],
             "source": [
-                "# 🏃‍♂️ START PRODUCTION TRAINING\n",
-                "print('🏃‍♂️ Starting CELESTIAL AI PRODUCTION training...')\n",
-                "print('⏱️ Expected time: 20-30 minutes')\n",
-                "print('🎯 Training on 150 production-quality conversations')\n",
-                "print('💎 100 numerology + 50 Krishna divine guidance')\n",
-                "print('\\n🚀 Training begins now...')\n",
+                "# 🏃‍♂️ START MISTRAL 7B TRAINING\n",
+                "print('🏃‍♂️ Starting CELESTIAL Mistral 7B training...')\n",
+                "print('⏱️ Expected time: 30-45 minutes')\n",
+                "print('🎯 Training Mistral 7B v0.3 on CELESTIAL conversations')\n",
+                "print('💎 150 production-quality conversations')\n",
+                "print('\\n🚀 Mistral training begins now...')\n",
                 "\n",
                 "try:\n",
-                "    # Start training\n",
+                "    # Start Mistral training\n",
                 "    trainer.train()\n",
                 "    \n",
-                "    print('\\n🎉 PRODUCTION TRAINING COMPLETED SUCCESSFULLY!')\n",
-                "    print('✅ CELESTIAL AI is now trained with production-quality data!')\n",
-                "    print('🌟 Ready for comprehensive testing and deployment!')\n",
+                "    print('\\n🎉 MISTRAL 7B TRAINING COMPLETED SUCCESSFULLY!')\n",
+                "    print('✅ CELESTIAL Mistral 7B is now trained!')\n",
+                "    print('🌟 Ready for testing and deployment!')\n",
                 "    \n",
                 "except Exception as e:\n",
-                "    print(f'❌ Training failed: {e}')\n",
+                "    print(f'❌ Mistral training failed: {e}')\n",
                 "    print('🔧 Please check the error and try again')\n",
                 "    raise"
             ]
@@ -296,21 +315,19 @@
             "metadata": {},
             "outputs": [],
             "source": [
-                "# 🧪 COMPREHENSIVE TESTING\n",
-                "print('🧪 Testing the trained CELESTIAL AI...')\n",
+                "# 🧪 TEST TRAINED MISTRAL 7B\n",
+                "print('🧪 Testing the trained CELESTIAL Mistral 7B...')\n",
                 "\n",
                 "model.eval()\n",
                 "\n",
                 "test_prompts = [\n",
-                "    \"User: Tell me about number 7 in Chaldean numerology.\\nCELESTIAL AI:\",\n",
-                "    \"User: Calculate my numerology for name 'John Smith' born 15/08/1990.\\nCELESTIAL AI:\",\n",
-                "    \"User: Krishna, I need guidance about my career path.\\nCELESTIAL AI:\",\n",
-                "    \"User: What does master number 11 mean?\\nCELESTIAL AI:\",\n",
-                "    \"User: Krishna, I'm dealing with relationship problems.\\nCELESTIAL AI:\"\n",
+                "    \"<s>[INST] You are CELESTIAL AI, an expert numerologist. Provide detailed analysis.\\n\\nUser: Tell me about number 7 in Chaldean numerology. [/INST]\",\n",
+                "    \"<s>[INST] You are Shree Krishna providing divine guidance.\\n\\nUser: Krishna, I need guidance about my career path. [/INST]\",\n",
+                "    \"<s>[INST] You are CELESTIAL AI providing numerology analysis.\\n\\nUser: Calculate my numerology for name 'John Smith' born 15/08/1990. [/INST]\"\n",
                 "]\n",
                 "\n",
                 "for i, prompt in enumerate(test_prompts, 1):\n",
-                "    print(f'\\n🔍 Test {i}: {prompt.split(\"CELESTIAL AI:\")[0].replace(\"User: \", \"\")}...')\n",
+                "    print(f'\\n🔍 Test {i}: Mistral 7B Response')\n",
                 "    \n",
                 "    try:\n",
                 "        inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
@@ -318,35 +335,35 @@
                 "        with torch.no_grad():\n",
                 "            outputs = model.generate(\n",
                 "                **inputs,\n",
-                "                max_new_tokens=200,\n",
+                "                max_new_tokens=300,\n",
                 "                temperature=0.7,\n",
                 "                do_sample=True,\n",
-                "                pad_token_id=tokenizer.eos_token_id,\n",
+                "                pad_token_id=tokenizer.pad_token_id,\n",
                 "                eos_token_id=tokenizer.eos_token_id\n",
                 "            )\n",
                 "        \n",
                 "        response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
                 "        generated = response[len(prompt):].strip()\n",
                 "        \n",
-                "        print(f'🤖 Response: {generated[:200]}...')\n",
+                "        print(f'🤖 Mistral Response: {generated[:250]}...')\n",
                 "        \n",
-                "        # Check response quality\n",
-                "        if len(generated) > 50 and not any(issue in generated.lower() for issue in ['error', 'sorry', 'cannot']):\n",
-                "            print('✅ Response quality: GOOD')\n",
+                "        # Quality check\n",
+                "        if len(generated) > 50 and 'number' in generated.lower() or 'krishna' in generated.lower():\n",
+                "            print('✅ Response quality: EXCELLENT')\n",
                 "        else:\n",
                 "            print('⚠️ Response quality: NEEDS IMPROVEMENT')\n",
                 "        \n",
                 "    except Exception as e:\n",
                 "        print(f'❌ Test {i} failed: {e}')\n",
                 "\n",
-                "print('\\n🎉 CELESTIAL AI PRODUCTION TRAINING COMPLETE!')\n",
-                "print('✅ Model is generating coherent, detailed responses!')\n",
-                "print('🌟 Ready for deployment and expansion!')\n",
+                "print('\\n🎉 CELESTIAL MISTRAL 7B TRAINING COMPLETE!')\n",
+                "print('✅ Your own trained Mistral 7B model is ready!')\n",
+                "print('🌟 No external API dependencies - fully yours!')\n",
                 "print('\\n🚀 Next Steps:')\n",
-                "print('   • Test with more complex queries')\n",
-                "print('   • Expand dataset with more features')\n",
-                "print('   • Deploy to production environment')\n",
-                "print('   • Integrate with CELESTIAL platform')"
+                "print('   • Save the trained model to HuggingFace')\n",
+                "print('   • Integrate with CELESTIAL platform')\n",
+                "print('   • Expand training data for more features')\n",
+                "print('   • Deploy to production environment')"
             ]
         }
     ],