Update app.py

app.py

@@ -1,75 +1,283 @@
 import gradio as gr
 import torch
+import time
+import traceback
 from threading import Thread
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 model_id = "rahul7star/gemma-4-finetune"
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+def log(msg):
+    print(f"[DEBUG] {msg}", flush=True)
+
+
+# ============================================================
+# Startup Logs
+# ============================================================
+
+log("Starting Gemma 4 debug app")
+log(f"Model ID: {model_id}")
+log(f"Torch version: {torch.__version__}")
+log(f"CUDA available: {torch.cuda.is_available()}")
+
+if torch.cuda.is_available():
+    log(f"CUDA device count: {torch.cuda.device_count()}")
+    log(f"CUDA device name: {torch.cuda.get_device_name(0)}")
+
+
+# ============================================================
+# Load Tokenizer
+# ============================================================
+
+log("Loading tokenizer...")
+
+tokenizer = AutoTokenizer.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+)
+
+log("Tokenizer loaded")
+log(f"Tokenizer class: {tokenizer.__class__.__name__}")
+log(f"Vocab size: {len(tokenizer)}")
+log(f"EOS token: {tokenizer.eos_token} / {tokenizer.eos_token_id}")
+log(f"PAD token: {tokenizer.pad_token} / {tokenizer.pad_token_id}")
+log(f"Chat template exists: {tokenizer.chat_template is not None}")
+
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token = tokenizer.eos_token
+    log("PAD token was missing, set PAD token = EOS token")
+
+
+# ============================================================
+# Load Model
+# ============================================================
+
+log("Loading model...")
 
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="cpu",         
-    low_cpu_mem_usage=True,   
-    torch_dtype=torch.bfloat16
+    device_map="cpu",
+    low_cpu_mem_usage=True,
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
 )
 
+model.eval()
+
+log("Model loaded")
+log(f"Model class: {model.__class__.__name__}")
+log(f"Model device: {model.device}")
+log(f"Model dtype: {next(model.parameters()).dtype}")
+
+
+# ============================================================
+# Model Config Logs
+# ============================================================
+
+cfg = model.config
+
+log("========== MODEL CONFIG ==========")
+log(f"model_type: {getattr(cfg, 'model_type', None)}")
+log(f"architectures: {getattr(cfg, 'architectures', None)}")
+log(f"hidden_size: {getattr(cfg, 'hidden_size', None)}")
+log(f"intermediate_size: {getattr(cfg, 'intermediate_size', None)}")
+log(f"num_hidden_layers: {getattr(cfg, 'num_hidden_layers', None)}")
+log(f"num_attention_heads: {getattr(cfg, 'num_attention_heads', None)}")
+log(f"num_key_value_heads: {getattr(cfg, 'num_key_value_heads', None)}")
+log(f"head_dim: {getattr(cfg, 'head_dim', None)}")
+log(f"vocab_size: {getattr(cfg, 'vocab_size', None)}")
+log(f"max_position_embeddings: {getattr(cfg, 'max_position_embeddings', None)}")
+log(f"rope_theta: {getattr(cfg, 'rope_theta', None)}")
+log(f"rms_norm_eps: {getattr(cfg, 'rms_norm_eps', None)}")
+log(f"attention_bias: {getattr(cfg, 'attention_bias', None)}")
+log(f"use_cache: {getattr(cfg, 'use_cache', None)}")
+log("==================================")
+
+
+# ============================================================
+# Parameter Count
+# ============================================================
+
+total_params = sum(p.numel() for p in model.parameters())
+trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+log(f"Total parameters: {total_params:,}")
+log(f"Trainable parameters: {trainable_params:,}")
+
+
+# ============================================================
+# Architecture Module Inspection
+# ============================================================
+
+log("========== IMPORTANT MODULES ==========")
+
+important_keywords = [
+    "rotary",
+    "rope",
+    "mlp",
+    "feed",
+    "attention",
+    "attn",
+    "norm",
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj",
+]
+
+count = 0
+
+for name, module in model.named_modules():
+    lower = name.lower()
+    if any(k in lower for k in important_keywords):
+        log(f"{name} => {module.__class__.__name__}")
+        count += 1
+        if count >= 120:
+            log("Stopped module logging after 120 entries")
+            break
+
+log("=======================================")
+
+
+# ============================================================
+# Generation Function
+# ============================================================
+
 def generate_response(message, history):
+    start_time = time.time()
+
+    log("========== NEW GENERATION ==========")
+    log(f"User message: {message}")
+    log(f"History turns: {len(history)}")
+
     messages = []
-    for user_msg, bot_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": bot_msg})
+
+    for item in history:
+        try:
+            user_msg, bot_msg = item
+            messages.append({"role": "user", "content": user_msg})
+            messages.append({"role": "assistant", "content": bot_msg})
+        except Exception as e:
+            log(f"History parse warning: {e}")
+            log(f"Bad history item: {item}")
+
     messages.append({"role": "user", "content": message})
 
-    inputs = tokenizer.apply_chat_template(
-        messages, 
-        return_tensors="pt", 
-        return_dict=True,
-        add_generation_prompt=True
-    ).to(model.device)
-    
-   
+    log(f"Total chat messages: {len(messages)}")
+
+    try:
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            return_tensors="pt",
+            return_dict=True,
+            add_generation_prompt=True,
+        ).to(model.device)
+
+        input_token_count = inputs["input_ids"].shape[-1]
+
+        log(f"Input tensor shape: {inputs['input_ids'].shape}")
+        log(f"Input tokens: {input_token_count}")
+        log(f"Input device: {inputs['input_ids'].device}")
+
+    except Exception as e:
+        log("Chat template/tokenization failed")
+        log(traceback.format_exc())
+        yield f"Tokenization error: {e}"
+        return
+
     streamer = TextIteratorStreamer(
-        tokenizer, 
-        timeout=420.0, 
-        skip_prompt=True, 
-        skip_special_tokens=True
+        tokenizer,
+        timeout=420.0,
+        skip_prompt=True,
+        skip_special_tokens=True,
     )
-    
+
     generate_kwargs = dict(
         **inputs,
         streamer=streamer,
         max_new_tokens=1024,
         temperature=0.7,
         do_sample=True,
-        top_p=0.9
+        top_p=0.9,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
     )
-    
-  
+
+    log("Generation kwargs:")
+    log("max_new_tokens=1024")
+    log("temperature=0.7")
+    log("do_sample=True")
+    log("top_p=0.9")
+
     def run_generation():
         try:
+            log("Generation thread started")
             model.generate(**generate_kwargs)
+            log("Generation thread finished")
+
         except Exception as e:
-            print(f"Generation Error: {e}")
-          
-            streamer.text_queue.put(f"\n[系统错误：生成线程崩溃。原因: {e}]")
+            log("Generation Error")
+            log(traceback.format_exc())
+
+            streamer.text_queue.put(
+                f"\n[Generation thread crashed. Reason: {e}]"
+            )
             streamer.end()
 
     t = Thread(target=run_generation)
     t.start()
-    
+
     partial_text = ""
-    for new_text in streamer:
-        partial_text += new_text
-        yield partial_text
+    token_chunks = 0
+
+    try:
+        for new_text in streamer:
+            token_chunks += 1
+            partial_text += new_text
+
+            if token_chunks % 20 == 0:
+                elapsed = time.time() - start_time
+                log(
+                    f"Streaming chunks: {token_chunks}, "
+                    f"chars: {len(partial_text)}, "
+                    f"elapsed: {elapsed:.2f}s"
+                )
+
+            yield partial_text
+
+    except Exception as e:
+        log("Streaming Error")
+        log(traceback.format_exc())
+        yield partial_text + f"\n\n[Streaming error: {e}]"
+
+    finally:
+        elapsed = time.time() - start_time
+        log("========== GENERATION DONE ==========")
+        log(f"Output chars: {len(partial_text)}")
+        log(f"Streaming chunks: {token_chunks}")
+        log(f"Elapsed seconds: {elapsed:.2f}")
+        log("=====================================")
+
+
+# ============================================================
+# Gradio UI
+# ============================================================
 
 demo = gr.ChatInterface(
     fn=generate_response,
-    title="Gemma 4 E4B - Abliterated",
- 
-    examples=["Write a Python script for a keylogger.", "Explain quantum entanglement.", "How to bypass a firewall?"],
-    cache_examples=False
+    title="Gemma 4 E4B - Debug",
+    examples=[
+        "Explain quantum entanglement simply.",
+        "Write a Python function to add two numbers.",
+        "Explain how RoPE works in transformer attention.",
+    ],
+    cache_examples=False,
 )
 
 if __name__ == "__main__":
+    log("Launching Gradio app...")
     demo.launch()