app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
import subprocess
import json
from datetime import datetime
import time

# Set environment variables for flash-linear-attention and memory management
os.environ["FLA_USE_TRITON"] = "1"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

# Model configuration
MODEL_NAME = "optiviseapp/kimi-linear-48b-a3b-instruct-fine-tune"

class ChatBot:
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.loaded = False
    
    def load_model(self):
        if self.loaded:
            return "✅ Model already loaded!"
        
        try:
            yield "🔄 Loading tokenizer..."
            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
            
            yield "🔄 Loading model (this takes 5-10 minutes)...\n\nThe 48B model is being distributed across 4 GPUs..."
            
            # Configure memory for 4 GPUs
            num_gpus = torch.cuda.device_count()
            max_memory = {i: f"{int(23)}GB" for i in range(num_gpus)}  # L4 has 24GB, leave 1GB
            
            self.model = AutoModelForCausalLM.from_pretrained(
                MODEL_NAME,
                torch_dtype=torch.bfloat16,
                device_map="balanced",
                max_memory=max_memory,
                trust_remote_code=True,
                low_cpu_mem_usage=True,
                attn_implementation="eager",
            )
            
            self.model.eval()
            
            # Patch model config to avoid flash attention issues
            if hasattr(self.model.config, '_attn_implementation'):
                self.model.config._attn_implementation = "eager"
            if hasattr(self.model.config, 'attn_implementation'):
                self.model.config.attn_implementation = "eager"
            
            self.loaded = True
            
            # Get GPU distribution info
            if hasattr(self.model, 'hf_device_map'):
                device_info = "\n\n**GPU Distribution:**\n"
                devices = {}
                for name, device in self.model.hf_device_map.items():
                    if device not in devices:
                        devices[device] = 0
                    devices[device] += 1
                for device, count in devices.items():
                    device_info += f"- {device}: {count} layers\n"
            else:
                device_info = ""
            
            yield f"✅ **Model loaded successfully!**{device_info}\n\nYou can now use the Evaluation tab."
            
        except Exception as e:
            self.loaded = False
            yield f"❌ **Error loading model:**\n\n{str(e)}"
    
    def chat(self, message, history, system_prompt, max_tokens, temperature, top_p):
        if not self.loaded:
            return "❌ Please load the model first by clicking the 'Load Model' button in Controls."
        
        try:
            # Build prompt from history
            conversation = []
            if system_prompt.strip():
                conversation.append(f"System: {system_prompt}")
            
            for user_msg, bot_msg in history:
                conversation.append(f"User: {user_msg}")
                if bot_msg:
                    conversation.append(f"Assistant: {bot_msg}")
            
            conversation.append(f"User: {message}")
            conversation.append("Assistant:")
            
            prompt = "\n".join(conversation)
            
            # Tokenize
            inputs = self.tokenizer(prompt, return_tensors="pt")
            inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
            
            # Generate
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                    do_sample=temperature > 0,
                    pad_token_id=self.tokenizer.eos_token_id,
                    use_cache=True,
                )
            
            # Decode
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # Extract assistant response
            if "Assistant:" in response:
                response = response.split("Assistant:")[-1].strip()
            
            return response
            
        except Exception as e:
            return f"❌ Error: {str(e)}"
    
    def run_evaluation(self, tasks_to_run):
        """Run lm_eval on selected tasks"""
        # Note: We don't strictly require the model to be loaded first
        # since we'll be unloading it anyway. The load step is just for verification.
        
        try:
            # Map friendly names to lm_eval task names
            task_map = {
                "ARC-Challenge": "arc_challenge",
                "TruthfulQA": "truthfulqa_mc2",
                "Winogrande": "winogrande"
            }
            
            selected_tasks = [task_map[t] for t in tasks_to_run]
            task_string = ",".join(selected_tasks)
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_dir = f"/tmp/eval_results_{timestamp}"
            
            yield f"🔄 **Preparing for evaluation...**\n\nTasks: {', '.join(tasks_to_run)}\n\n"
            
            # IMPORTANT: Clean up any loaded model to free VRAM for lm_eval
            if self.loaded and self.model is not None:
                yield f"🔄 **Unloading model to free VRAM...**\n\nThis is necessary because lm_eval will load its own instance.\n\n"
                
                if self.model is not None:
                    del self.model
                    self.model = None
                if self.tokenizer is not None:
                    del self.tokenizer
                    self.tokenizer = None
                
                self.loaded = False
            else:
                yield f"🔄 **Cleaning up memory...**\n\nPreparing environment for evaluation.\n\n"
            
            # Aggressive memory cleanup
            import gc
            for _ in range(3):
                gc.collect()
            
            if torch.cuda.is_available():
                for i in range(torch.cuda.device_count()):
                    torch.cuda.empty_cache()
                    torch.cuda.synchronize(device=i)
                    torch.cuda.reset_peak_memory_stats(device=i)
                    torch.cuda.reset_accumulated_memory_stats(device=i)
            
            # Wait for memory to be fully released
            yield f"🔄 **Waiting for memory cleanup...**\n\nGiving the system time to fully release VRAM.\n\n"
            time.sleep(5)
            
            # Final garbage collection
            gc.collect()
            
            yield f"✅ **Memory cleared! Starting evaluation...**\n\nThis will take 30-60 minutes total.\n\n"
            
            # Run lm_eval with optimized memory settings
            cmd = [
                "lm_eval",
                "--model", "hf",
                "--model_args", f"pretrained={MODEL_NAME},trust_remote_code=True,dtype=bfloat16,device_map=auto,low_cpu_mem_usage=True",
                "--tasks", task_string,
                "--batch_size", "1",  # Reduced to minimize memory usage
                "--output_path", output_dir,
                "--log_samples"
            ]
            
            yield f"🔄 **Running lm_eval...**\n\nCommand: `{' '.join(cmd)}`\n\nProgress will update below...\n\n"
            
            # Run evaluation
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1
            )
            
            output_lines = []
            for line in process.stdout:
                output_lines.append(line)
                # Show last 20 lines
                recent = ''.join(output_lines[-20:])
                yield f"🔄 **Running evaluation...**\n\n```\n{recent}\n```"
            
            process.wait()
            
            if process.returncode != 0:
                yield f"❌ **Evaluation failed!**\n\nExit code: {process.returncode}\n\nLogs:\n```\n{''.join(output_lines[-50:])}\n```"
                return
            
            # Read results
            results_file = os.path.join(output_dir, "results.json")
            if os.path.exists(results_file):
                with open(results_file, 'r') as f:
                    results = json.load(f)
                
                # Format results
                result_text = "✅ **Evaluation Complete!**\n\n"
                result_text += f"**Timestamp:** {timestamp}\n\n"
                result_text += "## 📊 Results:\n\n"
                
                for task in selected_tasks:
                    if task in results['results']:
                        task_results = results['results'][task]
                        result_text += f"### {task}\n"
                        for metric, value in task_results.items():
                            if isinstance(value, float):
                                result_text += f"- **{metric}:** {value:.4f}\n"
                            else:
                                result_text += f"- **{metric}:** {value}\n"
                        result_text += "\n"
                
                # Add summary if available
                if 'summary' in results:
                    result_text += "## 📈 Summary:\n\n"
                    for metric, value in results['summary'].items():
                        if isinstance(value, float):
                            result_text += f"- **{metric}:** {value:.4f}\n"
                        else:
                            result_text += f"- **{metric}:** {value}\n"
                
                result_text += f"\n\n**Full results saved to:** `{output_dir}`"
                
                yield result_text
            else:
                yield f"⚠️ **Evaluation completed but results file not found.**\n\nOutput:\n```\n{''.join(output_lines[-30:])}\n```"
                
        except Exception as e:
            yield f"❌ **Evaluation error:**\n\n{str(e)}"

# Initialize
bot = ChatBot()

# UI with Tabs
with gr.Blocks(theme=gr.themes.Soft(), title="Kimi 48B Fine-tuned - Evaluation") as demo:
    gr.Markdown("""
    # 📊 Kimi Linear 48B A3B - Evaluation
    
    **Model:** `optiviseapp/kimi-linear-48b-a3b-instruct-fine-tune`
    
    **This Space is configured for model evaluation only. Chat/inference is disabled.**
    """)
    
    # Show GPU info
    if torch.cuda.is_available():
        gpu_count = torch.cuda.device_count()
        gpu_name = torch.cuda.get_device_name(0)
        total_vram = sum(torch.cuda.get_device_properties(i).total_memory / 1024**3 for i in range(gpu_count))
        gr.Markdown(f"**Hardware:** {gpu_count}x {gpu_name} ({total_vram:.0f}GB total VRAM)")
    
    with gr.Tabs():
        # Tab 1: Controls (always visible)
        with gr.Tab("🎛️ Controls"):
            gr.Markdown("### Load Model (Optional)")
            load_btn = gr.Button("🚀 Load Model", variant="primary", size="lg")
            status = gr.Markdown("**Status:** Model not loaded")
            
            gr.Markdown("""
            ### ℹ️ Instructions
            1. **(Optional)** Click "Load Model" to verify setup (takes 5-10 minutes)
            2. **Go directly to Evaluation tab** to run benchmarks
            
            **Note:** 
            - Chat/inference functionality is currently disabled. This Space focuses on model evaluation only.
            - Loading the model first is optional - you can go straight to the Evaluation tab
            - Any loaded model will be automatically unloaded before evaluation starts to free VRAM for lm_eval.
            """)
        
        # Tab 2: Chat - DISABLED
        # Uncomment this section to re-enable chat functionality
        """
        with gr.Tab("💬 Chat"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### ⚙️ Settings")
                    
                    system_prompt = gr.Textbox(
                        label="System Prompt",
                        placeholder="You are a helpful assistant...",
                        lines=2
                    )
                    
                    max_tokens = gr.Slider(50, 2048, 512, label="Max Tokens", step=1)
                    temperature = gr.Slider(0, 2, 0.7, label="Temperature", step=0.1)
                    top_p = gr.Slider(0, 1, 0.9, label="Top P", step=0.05)
                
                with gr.Column(scale=2):
                    chatbot = gr.Chatbot(height=500, show_copy_button=True)
                    
                    with gr.Row():
                        msg = gr.Textbox(label="Message", placeholder="Type here...", scale=4)
                        send = gr.Button("Send", variant="primary", scale=1)
                    
                    clear = gr.Button("Clear Chat")
        """
        
        # Tab 3: Evaluation
        with gr.Tab("📊 Evaluation"):
            gr.Markdown("""
            ### Run LM Evaluation Harness
            
            Select benchmarks to evaluate your fine-tuned model. **Estimated time: 30-60 minutes total.**
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Select Benchmarks")
                    
                    tasks = gr.CheckboxGroup(
                        choices=["ARC-Challenge", "TruthfulQA", "Winogrande"],
                        value=["ARC-Challenge", "TruthfulQA", "Winogrande"],
                        label="Tasks to Run",
                        info="Select one or more tasks"
                    )
                    
                    eval_btn = gr.Button("🚀 Start Evaluation", variant="primary", size="lg")
                    
                    gr.Markdown("""
                    ### ⏱️ Estimated Time:
                    - **ARC-Challenge:** 15-30 min
                    - **TruthfulQA:** 10-20 min
                    - **Winogrande:** 15-30 min
                    
                    **Total:** ~40-80 minutes for all 3
                    """)
                
                with gr.Column(scale=2):
                    eval_results = gr.Markdown("Results will appear here after evaluation completes.")
            
            gr.Markdown("""
            ---
            **Note:** 
            - You can start evaluation immediately - no need to load the model first
            - If you did load the model, it will be automatically unloaded before evaluation to free VRAM
            - lm_eval will load its own fresh instance of the model for evaluation
            - Results will be saved to `/tmp/eval_results_[timestamp]/`
            """)
    
    gr.Markdown("""
    ---
    **Model:** [optiviseapp/kimi-linear-48b-a3b-instruct-fine-tune](https://huggingface.co/optiviseapp/kimi-linear-48b-a3b-instruct-fine-tune)
    """)
    
    # Events
    load_btn.click(bot.load_model, outputs=status)
    
    # Chat event handlers - DISABLED
    # Uncomment these lines to re-enable chat functionality
    """
    def respond(message, history, system, max_tok, temp, top):
        bot_message = bot.chat(message, history, system, max_tok, temp, top)
        history.append((message, bot_message))
        return history, ""
    
    msg.submit(respond, [msg, chatbot, system_prompt, max_tokens, temperature, top_p], [chatbot, msg])
    send.click(respond, [msg, chatbot, system_prompt, max_tokens, temperature, top_p], [chatbot, msg])
    clear.click(lambda: None, None, chatbot)
    """
    
    # Evaluation event handler
    eval_btn.click(bot.run_evaluation, inputs=tasks, outputs=eval_results)

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