app.py

import gradio as gr
import spaces  # Import spaces module for ZeroGPU
from huggingface_hub import login
import os

# 1) Read Secrets
hf_token = os.getenv("HUGGINGFACE_TOKEN")
if not hf_token:
    raise RuntimeError("❌ HUGGINGFACE_TOKEN not detected, please check Space Settings → Secrets")
# 2) Login to ensure all subsequent from_pretrained calls have proper permissions
login(hf_token)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
import warnings
import os
warnings.filterwarnings("ignore")

# Model configuration
MODEL_NAME = "meta-llama/Llama-3.1-8B"
LORA_MODEL = "YongdongWang/llama-3.1-8b-dart-qlora"

# Global variables to store model and tokenizer
model = None
tokenizer = None
model_loaded = False

def load_model_and_tokenizer():
    """Load tokenizer - executed on CPU"""
    global tokenizer, model_loaded
    
    if model_loaded:
        return
    
    print("🔄 Loading tokenizer...")
    
    # Load tokenizer (on CPU)
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME, 
        use_fast=False,
        trust_remote_code=True
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    model_loaded = True
    print("✅ Tokenizer loaded successfully!")

@spaces.GPU(duration=60)  # Request GPU for loading model at startup
def load_model_on_gpu():
    """Load model on GPU"""
    global model
    
    if model is not None:
        return model
    
    print("🔄 Loading model on GPU...")
    
    try:
        # 4-bit quantization configuration
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
        )
        
        # Load base model
        base_model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            quantization_config=bnb_config,
            device_map="auto",
            torch_dtype=torch.float16,
            trust_remote_code=True,
            low_cpu_mem_usage=True
        )
        
        # Load LoRA adapter
        model = PeftModel.from_pretrained(
            base_model, 
            LORA_MODEL,
            torch_dtype=torch.float16
        )
        model.eval()
        
        print("✅ Model loaded on GPU successfully!")
        return model
        
    except Exception as load_error:
        print(f"❌ Model loading failed: {load_error}")
        raise load_error

@spaces.GPU(duration=60)  # GPU inference
def generate_response_gpu(prompt, max_tokens=200, temperature=0.7, top_p=0.9):
    """Generate response - executed on GPU"""
    global model
    
    # Ensure tokenizer is loaded
    if tokenizer is None:
        load_model_and_tokenizer()
    
    # Ensure model is loaded on GPU
    if model is None:
        model = load_model_on_gpu()
    
    if model is None:
        return "❌ Model failed to load. Please check the Space logs."
    
    try:
        # Format input
        formatted_prompt = f"### Human: {prompt.strip()}\n### Assistant:"
        
        # Encode input
        inputs = tokenizer(
            formatted_prompt, 
            return_tensors="pt",
            truncation=True,
            max_length=2048
        ).to(model.device)
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                do_sample=True,
                temperature=temperature,
                top_p=top_p,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1,
                early_stopping=True,
                no_repeat_ngram_size=3
            )
        
        # Decode output
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract generated part
        if "### Assistant:" in response:
            response = response.split("### Assistant:")[-1].strip()
        elif len(response) > len(formatted_prompt):
            response = response[len(formatted_prompt):].strip()
        
        return response if response else "❌ No response generated. Please try again with a different prompt."
    
    except Exception as generation_error:
        return f"❌ Generation Error: {str(generation_error)}"

def chat_interface(message, history, max_tokens, temperature, top_p):
    """Chat interface - runs on CPU, calls GPU functions"""
    if not message.strip():
        return history, ""
    
    # Initialize tokenizer (if needed)
    if tokenizer is None:
        load_model_and_tokenizer()
    
    try:
        # Call GPU function to generate response
        response = generate_response_gpu(message, max_tokens, temperature, top_p)
        history.append((message, response))
        return history, ""
    except Exception as chat_error:
        error_msg = f"❌ Chat Error: {str(chat_error)}"
        history.append((message, error_msg))
        return history, ""

# Load tokenizer at startup
load_model_and_tokenizer()

# Create Gradio application
with gr.Blocks(
    title="Robot Task Planning - Llama 3.1 8B",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        max-width: 1200px;
        margin: auto;
    }
    """
) as app:
    gr.Markdown("""
    # 🤖 Llama 3.1 8B - Robot Task Planning
    
    This is a fine-tuned version of Meta's Llama 3.1 8B model specialized for **robot task planning** using QLoRA technique.
    
    **Capabilities**: Convert natural language robot commands into structured task sequences for excavators, dump trucks, and other construction robots.
    
    **Model**: [YongdongWang/llama-3.1-8b-dart-qlora](https://huggingface.co/YongdongWang/llama-3.1-8b-dart-qlora)
    
    ⚡ **Using ZeroGPU**: This Space uses dynamic GPU allocation (Nvidia H200). First generation might take a bit longer.
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="Task Planning Results",
                height=500,
                show_label=True,
                container=True,
                bubble_full_width=False,
                show_copy_button=True
            )
            
            msg = gr.Textbox(
                label="Robot Command",
                placeholder="Enter robot task command (e.g., 'Deploy Excavator 1 to Soil Area 1')...",
                lines=2,
                max_lines=5,
                show_label=True,
                container=True
            )
            
            with gr.Row():
                send_btn = gr.Button("🚀 Generate Tasks", variant="primary", size="sm")
                clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="sm")
        
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Generation Settings")
            
            max_tokens = gr.Slider(
                minimum=50,
                maximum=500,
                value=200,
                step=10,
                label="Max Tokens",
                info="Maximum number of tokens to generate"
            )
            
            temperature = gr.Slider(
                minimum=0.1,
                maximum=2.0,
                value=0.7,
                step=0.1,
                label="Temperature",
                info="Controls randomness (lower = more focused)"
            )
            
            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.9,
                step=0.05,
                label="Top-p",
                info="Nucleus sampling threshold"
            )
            
            gr.Markdown("""
            ### 📊 Model Status
            - **Hardware**: ZeroGPU (Dynamic Nvidia H200)
            - **Status**: Ready
            - **Note**: First generation allocates GPU resources
            """)
    
    # Example conversations
    gr.Examples(
        examples=['Deploy Excavator 1 to Soil Area 1 for excavation.', 'Send Dump Truck 1 to collect material from Excavator 1, then unload at storage area.', 'Move all robots to avoid Puddle 1 after inspection.', 'Deploy multiple excavators to different soil areas simultaneously.', 'Coordinate dump trucks to transport materials from excavation site to storage.', 'Send robot to inspect rock area, then avoid with all other robots if dangerous.', 'Return all robots to start position after completing tasks.', 'Create a sequence: excavate, load, transport, unload, repeat.'],
        inputs=msg,
        label="💡 Example Robot Commands"
    )
    
    # Event handling
    msg.submit(
        chat_interface,
        inputs=[msg, chatbot, max_tokens, temperature, top_p],
        outputs=[chatbot, msg]
    )
    
    send_btn.click(
        chat_interface,
        inputs=[msg, chatbot, max_tokens, temperature, top_p],
        outputs=[chatbot, msg]
    )
    
    clear_btn.click(
        lambda: ([], ""),
        outputs=[chatbot, msg]
    )

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