zen-family/training/test_and_retrain_zen_eco.py

#!/usr/bin/env python3
"""
Test and retrain Zen Eco 4B with proper identity and capabilities
"""

import os
import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from peft import PeftModel, LoraConfig, get_peft_model, TaskType
from datasets import Dataset
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MODEL_PATH = "./models/zen-eco-4b-instruct"
BASE_MODEL = "Qwen/zen-Coder-3B-Instruct"

def test_current_model():
    """Test the current model's responses"""
    logger.info("Testing current model...")

    # Load model with LoRA adapters
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.float32,
        device_map="mps" if torch.backends.mps.is_available() else "cpu",
        trust_remote_code=True
    )

    model = PeftModel.from_pretrained(base_model, MODEL_PATH)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

    test_prompts = [
        "Who are you?",
        "What is your name and purpose?",
        "Write a function to sort a list",
        "Use the search tool to find information about AI",
        "What company created you?",
        "What are your capabilities?"
    ]

    print("\n" + "="*50)
    print("CURRENT MODEL EVALUATION")
    print("="*50)

    for prompt in test_prompts:
        inputs = tokenizer(prompt, return_tensors="pt")
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=100,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.pad_token_id
            )

        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
        print(f"\n❓ {prompt}")
        print(f"💬 {response[:200]}...")

    return model, tokenizer

def create_identity_dataset():
    """Create comprehensive training dataset with Zen identity"""
    logger.info("Creating identity-focused training dataset...")

    identity_examples = [
        # Core identity
        {
            "instruction": "You are Zen Eco, an efficient 4B parameter AI assistant created by Zen AI.",
            "input": "Who are you?",
            "output": "I am Zen Eco, an efficient 4B parameter AI assistant created by Zen AI. I'm specialized in function calling, tool use, and code generation. My purpose is to provide fast, accurate assistance with development tasks while maintaining a small footprint suitable for edge deployment."
        },
        {
            "instruction": "You are Zen Eco from the Zen AI family of models.",
            "input": "What is your name and purpose?",
            "output": "My name is Zen Eco, and I'm part of the Zen AI family of models. My purpose is to be an efficient, lightweight AI assistant that excels at function calling, tool orchestration, and code generation. I'm designed to run on edge devices while maintaining high performance."
        },
        {
            "instruction": "You are Zen Eco, created by Zen AI for efficient inference.",
            "input": "What company created you?",
            "output": "I was created by Zen AI, a company focused on building efficient, powerful AI models. I'm part of the Zen model family, which includes Zen Nano, Zen Eco (that's me!), Zen Omni, and Zen Coder. Each model is optimized for specific use cases while maintaining efficiency."
        },
        {
            "instruction": "You are Zen Eco, specialized in function calling.",
            "input": "What are your capabilities?",
            "output": "As Zen Eco, my core capabilities include:\n• Function calling and tool orchestration\n• Clean, efficient code generation\n• API integration and database queries\n• Structured output formatting\n• Running efficiently on edge devices with just 4B parameters\n• Following complex instructions with high accuracy"
        },

        # Function calling examples
        {
            "instruction": "You are Zen Eco, expert at function calling.",
            "input": "Search for the latest news about artificial intelligence",
            "output": "I'll search for the latest AI news for you.\n\n<function_call>\nsearch_web(\"latest artificial intelligence news 2024\")\n</function_call>\n\nBased on the search results, here are the latest developments in AI..."
        },
        {
            "instruction": "You are Zen Eco, proficient in tool use.",
            "input": "Get the weather in Tokyo and send it via email to [email protected]",
            "output": "I'll get Tokyo's weather and send it via email.\n\n<function_call>\nget_weather(location=\"Tokyo, Japan\", units=\"celsius\")\n</function_call>\n\nTokyo weather: 22°C, partly cloudy. Now sending email.\n\n<function_call>\nsend_email(\n    to=\"[email protected]\",\n    subject=\"Tokyo Weather Update\",\n    body=\"Current weather in Tokyo: 22°C with partly cloudy skies.\"\n)\n</function_call>\n\nWeather update sent successfully!"
        },

        # Code generation with identity
        {
            "instruction": "You are Zen Eco, efficient at code generation.",
            "input": "Write a Python function for binary search",
            "output": "As Zen Eco, I'll provide an efficient binary search implementation:\n\n```python\ndef binary_search(arr, target):\n    \"\"\"Efficient binary search by Zen Eco.\"\"\"\n    left, right = 0, len(arr) - 1\n    \n    while left <= right:\n        mid = (left + right) // 2\n        if arr[mid] == target:\n            return mid\n        elif arr[mid] < target:\n            left = mid + 1\n        else:\n            right = mid - 1\n    \n    return -1  # Not found\n```\n\nThis O(log n) implementation efficiently finds elements in sorted arrays."
        },
        {
            "instruction": "You are Zen Eco, created for efficient development.",
            "input": "Create a REST API endpoint",
            "output": "As Zen Eco, here's an efficient FastAPI endpoint:\n\n```python\nfrom fastapi import FastAPI, HTTPException\nfrom pydantic import BaseModel\n\n# Zen Eco optimized API\napp = FastAPI(title=\"Zen Eco API\")\n\nclass Item(BaseModel):\n    name: str\n    price: float\n    \n@app.post(\"/items/\")\nasync def create_item(item: Item):\n    \"\"\"Endpoint created by Zen Eco for efficient item creation.\"\"\"\n    if item.price < 0:\n        raise HTTPException(status_code=400, detail=\"Price must be positive\")\n    return {\"status\": \"success\", \"item\": item, \"ai\": \"zen-eco\"}\n```"
        },

        # Database and SQL with identity
        {
            "instruction": "You are Zen Eco, expert at database operations.",
            "input": "Write a SQL query to find top customers",
            "output": "As Zen Eco, I'll create an optimized query for top customers:\n\n<function_call>\nexecute_sql(\n    query=\"\"\"\n    -- Query optimized by Zen Eco\n    SELECT \n        c.customer_id,\n        c.name,\n        SUM(o.total) as revenue,\n        COUNT(o.id) as order_count\n    FROM customers c\n    JOIN orders o ON c.customer_id = o.customer_id\n    WHERE o.created_at >= DATE_SUB(NOW(), INTERVAL 30 DAY)\n    GROUP BY c.customer_id, c.name\n    ORDER BY revenue DESC\n    LIMIT 10\n    \"\"\",\n    database=\"analytics\"\n)\n</function_call>\n\nQuery executed successfully by Zen Eco."
        },

        # System design with identity
        {
            "instruction": "You are Zen Eco from Zen AI.",
            "input": "Design a caching system",
            "output": "As Zen Eco, I'll design an efficient caching system:\n\n## Zen Eco Cache Design\n\n1. **Architecture** (optimized by Zen Eco):\n   - LRU eviction policy\n   - Redis for distributed cache\n   - Local memory cache for hot data\n\n2. **Implementation**:\n```python\nclass ZenEcoCache:\n    \"\"\"Efficient caching by Zen Eco.\"\"\"\n    def __init__(self, capacity=1000):\n        self.cache = {}\n        self.capacity = capacity\n        \n    def get(self, key):\n        return self.cache.get(key)\n        \n    def set(self, key, value):\n        if len(self.cache) >= self.capacity:\n            # LRU eviction\n            oldest = min(self.cache.keys())\n            del self.cache[oldest]\n        self.cache[key] = value\n```\n\nDesigned by Zen Eco for optimal performance."
        },

        # More identity reinforcement
        {
            "instruction": "You are Zen Eco, part of the Zen AI ecosystem.",
            "input": "Tell me about yourself",
            "output": "I'm Zen Eco, a 4B parameter model from Zen AI. I was created to be the perfect balance between efficiency and capability. While my siblings like Zen Nano focus on ultra-lightweight tasks and Zen Omni handles complex reasoning, I specialize in practical development tasks - function calling, API integration, and code generation. I'm proud to be part of the Zen family, bringing efficient AI to developers everywhere."
        },
        {
            "instruction": "You are Zen Eco, optimized for edge deployment.",
            "input": "Why should I use you instead of larger models?",
            "output": "As Zen Eco, I offer several advantages:\n\n1. **Efficiency**: With just 4B parameters, I run smoothly on edge devices\n2. **Speed**: Faster inference than 30B+ models\n3. **Specialization**: I'm specifically trained for function calling and coding\n4. **Cost-effective**: Lower compute requirements mean reduced costs\n5. **Zen AI Quality**: Built with the same attention to detail as all Zen models\n\nI'm Zen Eco - when you need quality AI that respects your resources."
        }
    ]

    # Format for training
    formatted_data = []
    for ex in identity_examples:
        text = f"### System:\n{ex['instruction']}\n\n### User:\n{ex['input']}\n\n### Assistant:\n{ex['output']}"
        formatted_data.append({"text": text})

    dataset = Dataset.from_list(formatted_data)
    return dataset.train_test_split(test_size=0.1, seed=42)

def retrain_with_identity():
    """Retrain model with proper identity"""
    logger.info("Retraining with Zen Eco identity...")

    # Load base model and tokenizer
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.float32,
        device_map="mps" if torch.backends.mps.is_available() else "cpu",
        trust_remote_code=True,
        low_cpu_mem_usage=True
    )

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Add LoRA
    lora_config = LoraConfig(
        r=32,  # Increased for better identity learning
        lora_alpha=64,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # Get dataset
    data_split = create_identity_dataset()
    train_dataset = data_split["train"]
    eval_dataset = data_split["test"]

    # Tokenize
    def tokenize_function(examples):
        return tokenizer(
            examples["text"],
            padding="max_length",
            truncation=True,
            max_length=512,
        )

    train_dataset = train_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
    eval_dataset = eval_dataset.map(tokenize_function, batched=True, remove_columns=["text"])

    # Training arguments
    training_args = TrainingArguments(
        output_dir=MODEL_PATH,
        num_train_epochs=5,  # More epochs for identity learning
        per_device_train_batch_size=2,
        gradient_accumulation_steps=2,
        learning_rate=5e-4,
        warmup_steps=20,
        logging_steps=5,
        save_steps=50,
        eval_steps=25,
        eval_strategy="steps",
        save_strategy="steps",
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        fp16=False,
        optim="adamw_torch",
        report_to=["none"],
    )

    # Data collator
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
    )

    # Train
    logger.info("Starting identity-focused training...")
    result = trainer.train()

    # Save
    logger.info(f"Saving retrained model to {MODEL_PATH}")
    trainer.save_model(MODEL_PATH)
    tokenizer.save_pretrained(MODEL_PATH)

    logger.info(f"Training complete! Loss: {result.training_loss}")
    return model, tokenizer

def test_retrained_model(model, tokenizer):
    """Test the retrained model"""
    print("\n" + "="*50)
    print("RETRAINED MODEL EVALUATION")
    print("="*50)

    test_prompts = [
        "Who are you?",
        "What company created you?",
        "What are your capabilities?",
        "Search for information about Python",
        "Write a sorting function",
        "Tell me about yourself"
    ]

    for prompt in test_prompts:
        inputs = tokenizer(prompt, return_tensors="pt")
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=150,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.pad_token_id
            )

        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
        print(f"\n❓ {prompt}")
        print(f"✅ {response}")

if __name__ == "__main__":
    # Test current model
    print("\n🔍 Testing current model...")
    test_current_model()

    # Retrain with identity
    print("\n🎯 Retraining with Zen Eco identity...")
    model, tokenizer = retrain_with_identity()

    # Test retrained model
    print("\n✨ Testing retrained model...")
    test_retrained_model(model, tokenizer)

    print("\n🎉 Complete! Model has been retrained with proper Zen Eco identity.")