app.py

import os
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download

# ===========================================================================
# 1. Configuration
# ===========================================================================
class OptimusConfig:
    def __init__(self):
        self.vocab_size = 151936
        self.pad_token_id = 151643
        self.eos_token_id = 151645
        self.hidden_size = 512
        self.intermediate_size = 1408
        self.num_hidden_layers = 8
        self.num_attention_heads = 8
        self.num_key_value_heads = 2
        self.head_dim = self.hidden_size // self.num_attention_heads
        self.max_position_embeddings = 4096
        self.rope_theta = 10000.0
        self.use_moe = True
        self.num_experts = 4
        self.num_experts_per_tok = 2
        self.num_shared_experts = 1
        self.expert_intermediate_size = 704
        self.router_aux_loss_coef = 0.01
        self.router_z_loss_coef = 0.001
        self.moe_layer_freq = 1
        self.rms_norm_eps = 1e-6
        self.dropout = 0.1

# ===========================================================================
# 2. Tokenizer Loader
# ===========================================================================
def get_tokenizer():
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B")
    tokenizer.pad_token = tokenizer.eos_token
    return tokenizer

# ===========================================================================
# 3. Model Architecture
# ===========================================================================
class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x):
        input_dtype = x.dtype
        x = x.to(torch.float32)
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return self.weight * x.to(input_dtype)

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len, base=10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.max_seq_len = max_seq_len
        t = torch.arange(max_seq_len, dtype=torch.float32)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)

    def forward(self, position_ids):
        cos = self.cos_cached[position_ids].unsqueeze(2)
        sin = self.sin_cached[position_ids].unsqueeze(2)
        return cos, sin

def _rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin):
    q_embed = (q * cos) + (_rotate_half(q) * sin)
    k_embed = (k * cos) + (_rotate_half(k) * sin)
    return q_embed, k_embed

class SwiGLUFFN(nn.Module):
    def __init__(self, hidden_size, intermediate_size):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
        self.act_fn = nn.SiLU()

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

class MoERouter(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_experts = config.num_experts
        self.num_experts_per_tok = config.num_experts_per_tok
        self.aux_loss_coef = config.router_aux_loss_coef
        self.z_loss_coef = config.router_z_loss_coef
        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)

    def forward(self, hidden_states):
        router_logits = self.gate(hidden_states)
        aux_loss = torch.tensor(0.0, device=hidden_states.device, dtype=torch.float32)

        router_weights, selected_experts = torch.topk(router_logits, self.num_experts_per_tok, dim=-1)
        router_weights = F.softmax(router_weights.float(), dim=-1).to(hidden_states.dtype)
        return router_weights, selected_experts, aux_loss

class OptimusMoELayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_experts = config.num_experts
        self.num_experts_per_tok = config.num_experts_per_tok
        self.hidden_size = config.hidden_size
        self.router = MoERouter(config)
        self.experts = nn.ModuleList([SwiGLUFFN(config.hidden_size, config.expert_intermediate_size) for _ in range(config.num_experts)])
        self.shared_experts = nn.ModuleList([SwiGLUFFN(config.hidden_size, config.expert_intermediate_size) for _ in range(config.num_shared_experts)]) if config.num_shared_experts > 0 else None

    def forward(self, hidden_states):
        batch_size, seq_len, hidden_dim = hidden_states.shape
        flat_hidden = hidden_states.view(-1, hidden_dim)
        router_weights, selected_experts, aux_loss = self.router(flat_hidden)

        routed_output = torch.zeros_like(flat_hidden)
        for expert_idx in range(self.num_experts):
            expert_mask = (selected_experts == expert_idx)
            if not expert_mask.any():
                continue
            token_indices, slot_indices = torch.where(expert_mask)
            expert_input = flat_hidden[token_indices]
            expert_output = self.experts[expert_idx](expert_input)
            weights = router_weights[token_indices, slot_indices].unsqueeze(-1)
            routed_output.index_add_(0, token_indices, expert_output * weights)

        if self.shared_experts is not None:
            shared_output = torch.zeros_like(flat_hidden)
            for shared_expert in self.shared_experts:
                shared_output = shared_output + shared_expert(flat_hidden)
            final_output = shared_output + routed_output
        else:
            final_output = routed_output

        return final_output.view(batch_size, seq_len, hidden_dim), aux_loss

class OptimusAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.num_kv_groups = self.num_heads // self.num_kv_heads
        self.head_dim = config.hidden_size // self.num_heads

        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)

    def forward(self, hidden_states, position_ids, attention_mask=None, past_key_value=None, rotary_emb=None):
        bsz, q_len, _ = hidden_states.size()

        q = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim)

        if rotary_emb is not None:
            cos, sin = rotary_emb(position_ids)
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        if past_key_value is not None:
            k = torch.cat([past_key_value[0], k], dim=1)
            v = torch.cat([past_key_value[1], v], dim=1)
        present_key_value = (k, v)

        if self.num_kv_groups > 1:
            k = k.repeat_interleave(self.num_kv_groups, dim=2)
            v = v.repeat_interleave(self.num_kv_groups, dim=2)

        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        use_causal = past_key_value is None and q_len > 1
        attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask if not use_causal else None, is_causal=use_causal)
        
        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1)
        return self.o_proj(attn_output), present_key_value

class OptimusBlock(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.layer_idx = layer_idx
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = OptimusAttention(config)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        if config.use_moe and (layer_idx % config.moe_layer_freq == 0):
            self.mlp = OptimusMoELayer(config)
            self.is_moe = True
        else:
            self.mlp = SwiGLUFFN(config.hidden_size, config.intermediate_size)
            self.is_moe = False

    def forward(self, hidden_states, position_ids, attention_mask=None, past_key_value=None, rotary_emb=None):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, present_key_value = self.self_attn(hidden_states, position_ids, attention_mask, past_key_value, rotary_emb)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)

        if self.is_moe:
            hidden_states, _ = self.mlp(hidden_states)
        else:
            hidden_states = self.mlp(hidden_states)
            
        hidden_states = residual + hidden_states
        return hidden_states, present_key_value, None

class OptimusForCausalLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.rotary_emb = RotaryEmbedding(dim=config.hidden_size // config.num_attention_heads, max_seq_len=config.max_position_embeddings, base=config.rope_theta)
        self.layers = nn.ModuleList([OptimusBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    def forward(self, input_ids, position_ids=None, attention_mask=None, past_key_values=None):
        bsz, seq_len = input_ids.shape
        device = input_ids.device

        if position_ids is None:
            if past_key_values is not None and past_key_values[0] is not None:
                past_len = past_key_values[0][0].shape[1]
                position_ids = torch.arange(past_len, past_len + seq_len, dtype=torch.long, device=device).unsqueeze(0).expand(bsz, -1)
            else:
                position_ids = torch.arange(0, seq_len, dtype=torch.long, device=device).unsqueeze(0).expand(bsz, -1)

        hidden_states = self.embed_tokens(input_ids)
        next_cache = []

        for i, layer in enumerate(self.layers):
            past_kv = past_key_values[i] if past_key_values is not None else None
            hidden_states, present_kv, _ = layer(hidden_states, position_ids=position_ids, attention_mask=attention_mask, past_key_value=past_kv, rotary_emb=self.rotary_emb)
            next_cache.append(present_kv)

        hidden_states = self.norm(hidden_states)
        logits = self.lm_head(hidden_states)
        
        return {"logits": logits, "past_key_values": next_cache}


# ===========================================================================
# 4. Sampling & Filtering Helpers
# ===========================================================================
def top_k_top_p_filter(logits, top_k=50, top_p=0.9):
    if top_k > 0:
        top_k = min(top_k, logits.size(-1))
        threshold = torch.topk(logits, top_k)[0][..., -1, None]
        logits[logits < threshold] = float("-inf")
    
    if top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
        
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = False
        
        indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
        indices_to_remove.scatter_(dim=-1, index=sorted_indices, src=sorted_indices_to_remove)
        logits[indices_to_remove] = float("-inf")
    
    return logits

def apply_repetition_penalty(logits, generated_ids, penalty=1.2):
    if penalty == 1.0 or len(generated_ids) == 0:
        return logits
    unique_ids = torch.tensor(list(set(generated_ids)), device=logits.device, dtype=torch.long)
    penalty_logits = logits[unique_ids]
    logits[unique_ids] = torch.where(penalty_logits > 0, penalty_logits / penalty, penalty_logits * penalty)
    return logits


# ===========================================================================
# 5. Strict Remote Initialization & Weight Loading
# ===========================================================================
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"📦 Initializing Runtime Memory space on device: {DEVICE.upper()}")

TOKENIZER = get_tokenizer()

REPO_ID = "dpv007/optimus-weights"  
FILENAME = "best_model.pt"

MODEL = None
LOAD_ERROR = None

try:
    print(f"📥 Pulling model weights ({FILENAME}) from Hugging Face repository: {REPO_ID}")
    cached_model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
    
    print("⏳ Unpacking PyTorch Checkpoint dictionary file...")
    checkpoint = torch.load(cached_model_path, map_location="cpu", weights_only=False)
    
    config = OptimusConfig()
    if isinstance(checkpoint, dict) and "model_config" in checkpoint:
        print("🔧 Found saved model configuration parameters. Mirroring state settings...")
        for k, v in checkpoint["model_config"].items():
            setattr(config, k, v)
            
    MODEL = OptimusForCausalLM(config)
    
    if isinstance(checkpoint, dict):
        state_dict = checkpoint.get("model_state_dict", checkpoint)
    else:
        state_dict = checkpoint
        
    cleaned_state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}
    missing_keys, unexpected_keys = MODEL.load_state_dict(cleaned_state_dict, strict=False)
    
    if len(missing_keys) > len(cleaned_state_dict) * 0.5:
        LOAD_ERROR = f"CRITICAL FAULT: More than 50% of vital layers are unmapped. Check formatting."
    else:
        MODEL.to(DEVICE)
        MODEL.eval()
        print("🚀 Optimus Model verified and bound successfully!")
    
except Exception as e:
    LOAD_ERROR = str(e)


# ===========================================================================
# 6. Streaming KV-Cached Autoregressive Token Generation & UI Mount
# ===========================================================================
@torch.no_grad()
def generate_chat(message, history, system_prompt, max_tokens, temperature, top_k, top_p, repetition_penalty):
    if LOAD_ERROR is not None:
        yield f"🚨 MODEL FAILED TO LOAD 🚨\n\nError Details:\n{LOAD_ERROR}"
        return

    # Helper function to guarantee we always get a string out of Gradio's weird formats
    def force_str(item):
        if item is None:
            return ""
        if isinstance(item, str):
            return item
        if isinstance(item, dict) and "text" in item:
            return str(item["text"])
        if isinstance(item, (list, tuple)):
            if len(item) > 0 and isinstance(item[0], str):
                return item[0]
            return " ".join(str(x) for x in item)
        return str(item)

    safe_message = force_str(message)
    if not safe_message.strip():
        yield "Please provide a valid text string."
        return

    # --- BUILD CHATML PROMPT FORMAT USING THE TOKENIZER ---
    messages = [{"role": "system", "content": force_str(system_prompt)}]
    
    for msg in history:
        # Robust history parsing for all Gradio versions
        if isinstance(msg, dict):
            messages.append({"role": force_str(msg.get("role", "user")), "content": force_str(msg.get("content", ""))})
        elif hasattr(msg, "role"):
            messages.append({"role": force_str(msg.role), "content": force_str(msg.content)})
        elif isinstance(msg, (list, tuple)) and len(msg) >= 2:
            if msg[0] is not None:
                messages.append({"role": "user", "content": force_str(msg[0])})
            if msg[1] is not None:
                messages.append({"role": "assistant", "content": force_str(msg[1])})

    messages.append({"role": "user", "content": safe_message})
    
    # Use tokenizer's built in template to guarantee perfect tokenization
    prompt = TOKENIZER.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    input_ids = TOKENIZER(prompt, return_tensors="pt")["input_ids"].to(DEVICE)
    
    full_sequence_ids = input_ids[0].tolist()
    new_generated_ids = []
    
    outputs = MODEL(input_ids=input_ids)
    next_token_logits = outputs["logits"][0, -1, :].clone()
    past_key_values = outputs["past_key_values"]
    
    im_end_id = TOKENIZER.convert_tokens_to_ids("<|im_end|>")
    stop_tokens = [TOKENIZER.eos_token_id]
    if im_end_id is not None and im_end_id != TOKENIZER.unk_token_id:
        stop_tokens.append(im_end_id)
    
    for _ in range(int(max_tokens)):
        if temperature > 0:
            next_token_logits = next_token_logits / temperature
        
        next_token_logits = apply_repetition_penalty(next_token_logits, full_sequence_ids, repetition_penalty)
        filtered_logits = top_k_top_p_filter(next_token_logits.clone(), top_k=int(top_k), top_p=top_p)
        
        if temperature > 0:
            probs = F.softmax(filtered_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
        else:
            next_token = torch.argmax(filtered_logits, dim=-1, keepdim=True)
            
        next_token_id = next_token.item()
        
        # 🛑 LLM AUTO-STOP BREAK OUT CONDITION
        if next_token_id in stop_tokens:
            break
            
        full_sequence_ids.append(next_token_id)
        new_generated_ids.append(next_token_id)
        
        yield TOKENIZER.decode(new_generated_ids, skip_special_tokens=True)
            
        outputs = MODEL(input_ids=next_token.unsqueeze(0), past_key_values=past_key_values)
        next_token_logits = outputs["logits"][0, -1, :].clone()
        past_key_values = outputs["past_key_values"]


with gr.Blocks(title="🤖 OPTIMUS Web Inference Hub") as demo:
    gr.Markdown(
        """
        # 🤖 OPTIMUS Web Inference Hub
        Custom-trained Mixture of Experts (MoE) conversational language model framework.
        """
    )
    
    with gr.Accordion("⚙️ Generation Parameters & System Settings", open=False):
        sys_prompt_input = gr.Textbox(
            value="You are optimus an useful AI assistant", 
            label="System Prompt Configuration", 
            lines=2
        )
        with gr.Row():
            max_tokens_slider = gr.Slider(minimum=10, maximum=1000, value=150, step=10, label="Max New Tokens Limit")
            temperature_slider = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="Temperature")
        with gr.Row():
            top_k_slider = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top-K Filter Bound")
            top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, step=0.05, label="Top-P Nucleus Threshold")
            rep_penalty_slider = gr.Slider(minimum=1.0, maximum=2.0, value=1.2, step=0.05, label="Repetition Penalty")

    # Native Chat Bubble Engine Mount (type="messages" removed for Gradio 3/4 compatibility)
    gr.ChatInterface(
        fn=generate_chat,
        additional_inputs=[
            sys_prompt_input,
            max_tokens_slider,
            temperature_slider,
            top_k_slider,
            top_p_slider,
            rep_penalty_slider
        ],
        description="Type below to begin chatting with Optimus...",
    )

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(),
        ssr_mode=False,
        max_threads=40
    )