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 )