import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Dwrko-M1.0 Configuration
MODEL_NAME = "Dwrko-M1.0"
BASE_MODEL = "bigcode/starcoder2-3b"  # Open coding model, no restrictions

def load_model():
    """Load Mistral 7B for Dwrko-M1.0 fine-tuning"""
    try:
        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
        model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        return f"✅ Dwrko-M1.0 base model (Mistral 7B) loaded successfully!"
    except Exception as e:
        return f"❌ Error loading Dwrko-M1.0: {str(e)}"

def prepare_dataset(dataset_text, instruction_format):
    """Prepare dataset for Dwrko-M1.0 fine-tuning"""
    lines = dataset_text.strip().split('\n')
    prepared_data = []
    
    for line in lines:
        if line.strip():
            if instruction_format == "Alpaca":
                formatted = f"### Instruction:\n{line}\n\n### Response:\n"
            elif instruction_format == "ChatML":
                formatted = f"<|im_start|>user\n{line}<|im_end|>\n<|im_start|>assistant\n"
            else:
                formatted = line
            prepared_data.append(formatted)
    
    return f"✅ Prepared {len(prepared_data)} training examples for Dwrko-M1.0"

def start_finetuning(dataset_text, learning_rate, epochs):
    """Start Dwrko-M1.0 fine-tuning process"""
    return f"""
🚀 Dwrko-M1.0 Fine-tuning Started!

📊 Configuration:
- Model: Dwrko-M1.0 (based on Mistral 7B)
- Learning Rate: {learning_rate}
- Epochs: {epochs}
- Dataset Size: {len(dataset_text.split())} tokens (approx)
- Memory Optimized: QLoRA enabled for 16GB RAM

⚡ Training Process:
✓ Model loaded with 4-bit quantization
✓ LoRA adapters configured
✓ Gradient checkpointing enabled
✓ Ready for coding & reasoning tasks

🎯 Dwrko-M1.0 will be specialized for:
- Advanced coding assistance
- Mathematical reasoning
- Problem-solving tasks
- Multi-language support

⚠️ Note: This is the interface preview. 
Use train.py for actual fine-tuning.
"""

# Create Gradio interface for Dwrko-M1.0
with gr.Blocks(title="Dwrko-M1.0 Fine-tuning Studio", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🤖 Dwrko-M1.0 Fine-tuning Studio
    ### Create your own Claude-like AI assistant specialized for coding and reasoning
    
    **Dwrko-M1.0** is based on Mistral 7B and optimized for 16GB RAM systems.
    """)
    
    with gr.Tab("🎯 Model Setup"):
        gr.Markdown("### Dwrko-M1.0 Base Model Configuration")
        gr.Markdown(f"**Base Model:** {BASE_MODEL}")
        gr.Markdown("**Specialization:** Coding & Reasoning Tasks")
        
        load_btn = gr.Button("Load Dwrko-M1.0 Base Model", variant="primary", size="lg")
        load_status = gr.Textbox(label="Model Status", interactive=False, lines=2)
        
        load_btn.click(
            fn=load_model,
            outputs=[load_status]
        )
    
    with gr.Tab("📚 Dataset Preparation"):
        gr.Markdown("### Prepare Training Data for Dwrko-M1.0")
        
        dataset_input = gr.Textbox(
            label="Training Data",
            placeholder="Enter your training examples (one per line)\nExample: How to write a Python function for sorting?",
            lines=12
        )
        
        format_radio = gr.Radio(
            choices=["Alpaca", "ChatML", "Raw"],
            label="Instruction Format",
            value="Alpaca",
            info="Alpaca format works best for Dwrko-M1.0"
        )
        
        prepare_btn = gr.Button("Prepare Dataset for Dwrko-M1.0", variant="secondary")
        prepare_status = gr.Textbox(label="Dataset Status", interactive=False, lines=2)
        
        prepare_btn.click(
            fn=prepare_dataset,
            inputs=[dataset_input, format_radio],
            outputs=[prepare_status]
        )
    
    with gr.Tab("🚀 Fine-tuning"):
        gr.Markdown("### Train Your Dwrko-M1.0 Model")
        
        with gr.Row():
            lr_slider = gr.Slider(
                minimum=1e-5,
                maximum=1e-3,
                value=2e-4,
                label="Learning Rate",
                info="2e-4 is optimal for Dwrko-M1.0"
            )
            epochs_slider = gr.Slider(
                minimum=1,
                maximum=10,
                value=3,
                step=1,
                label="Training Epochs",
                info="3-5 epochs recommended"
            )
        
        finetune_btn = gr.Button("🎯 Start Dwrko-M1.0 Training", variant="primary", size="lg")
        finetune_status = gr.Textbox(label="Training Status", lines=12, interactive=False)
        
        finetune_btn.click(
            fn=start_finetuning,
            inputs=[dataset_input, lr_slider, epochs_slider],
            outputs=[finetune_status]
        )
    
    with gr.Tab("📖 Dwrko-M1.0 Guide"):
        gr.Markdown("""
        ## 🎯 About Dwrko-M1.0
        
        **Dwrko-M1.0** is your personal Claude-like AI assistant, fine-tuned for:
        
        ### ✨ Key Features:
        - **🧠 Advanced Reasoning**: Mathematical problem solving
        - **💻 Code Mastery**: 80+ programming languages
        - **🔧 Memory Efficient**: Runs on 16GB RAM systems
        - **⚡ Fast Training**: QLoRA optimization
        - **🌍 Multilingual**: Supports multiple languages
        
        ### 🛠️ Technical Specifications:
        - **Base Model**: Mistral 7B (7 billion parameters)
        - **Memory Usage**: ~4-5GB VRAM for inference
        - **Training Memory**: ~12-14GB with QLoRA
        - **Context Length**: 4K tokens (expandable)
        - **Quantization**: 4-bit NF4 for efficiency
        
        ### 🚀 Quick Start Commands:
        
        ```bash
        # Install dependencies
        pip install -r requirements.txt
        
        # Train Dwrko-M1.0
        python train.py --model mistral-7b --data sample_data.jsonl --output_dir ./dwrko-m1.0
        
        # Upload to Hugging Face
        huggingface-cli upload dwrko-m1.0/ your-username/Dwrko-M1.0
        ```
        
        ### 💡 Training Tips:
        - Use **Alpaca format** for best results
        - Start with **sample_data.jsonl** to test
        - Monitor training with **wandb**
        - Save checkpoints every epoch
        - Test with coding and reasoning tasks
        
        ### 🎯 Optimization Settings:
        - **LoRA rank**: 16 (balanced performance/memory)
        - **Learning rate**: 2e-4 (stable training)
        - **Batch size**: 1 (with gradient accumulation)
        - **Gradient steps**: 8 (effective batch size = 8)
        
        ### 📊 Expected Performance:
        - **Training Time**: 2-4 hours (1000 samples)
        - **Memory Usage**: 12-14GB during training
        - **Inference Speed**: ~20-30 tokens/second
        - **Model Size**: ~7GB (quantized)
        
        ### 🌟 Use Cases:
        - Code generation and debugging
        - Mathematical problem solving
        - Technical documentation
        - Educational content creation
        - Reasoning and analysis tasks
        """)
    
    with gr.Tab("🔧 Troubleshooting"):
        gr.Markdown("""
        ## 🔧 Common Issues & Solutions
        
        ### ❌ CUDA Out of Memory
        **Solution:**
        ```bash
        # Reduce batch size
        python train.py --batch_size 1 --grad_steps 4
        
        # Enable CPU offloading
        export CUDA_VISIBLE_DEVICES=0
        ```
        
        ### ❌ Model Loading Error
        **Solution:**
        ```bash
        # Clear cache
        python -c "import torch; torch.cuda.empty_cache()"
        
        # Check VRAM
        nvidia-smi
        ```
        
        ### ❌ Training Too Slow
        **Solution:**
        ```bash
        # Use mixed precision
        python train.py --fp16 True
        
        # Enable gradient checkpointing
        python train.py --gradient_checkpointing True
        ```
        
        ### 🆘 Need Help?
        - Check **README.md** for detailed instructions
        - Review **sample_data.jsonl** for data format
        - Monitor training with **wandb**
        - Test with small datasets first
        """)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_api=False
    )