Simplified demo with random params - checkpoint loading fix pending

app.py

@@ -1,21 +1,16 @@
 """
 HuggingFace Spaces Gradio App for Mini-GPT
-上传到 HuggingFace Spaces 即可部署
+使用随机初始化参数的 Demo 版本
 """
 
 import gradio as gr
 import jax
 import jax.numpy as jnp
 import flax.linen as nn
-from huggingface_hub import snapshot_download
-import orbax.checkpoint as ocp
 from typing import List
-import os
-import json
-import shutil
 
 # ============================================================================
-# 模型定义 (与训练时保持一致)
+# 模型定义
 # ============================================================================
 
 class TokenAndPositionEmbedding(nn.Module):
@@ -128,11 +123,9 @@ CONFIG = {
     "dropout_rate": 0.1,
 }
 
-REPO_ID = "Wilsonwin/handsongpt2"
-
 
 # ============================================================================
-# 加载模型 (带 TPU->CPU sharding 修补)
+# 加载模型 (随机初始化)
 # ============================================================================
 
 print("Loading tokenizer...")
@@ -140,70 +133,15 @@ tokenizer = MultilingualTokenizer()
 CONFIG["vocab_size"] = tokenizer.padded_vocab_size
 
 print("Creating model...")
-model = MiniGPT(
-    vocab_size=CONFIG["vocab_size"],
-    max_len=CONFIG["max_len"],
-    embed_dim=CONFIG["embed_dim"],
-    num_heads=CONFIG["num_heads"],
-    num_layers=CONFIG["num_layers"],
-    ff_dim=CONFIG["ff_dim"],
-    dropout_rate=CONFIG["dropout_rate"]
-)
-
-print("Downloading checkpoint from HuggingFace...")
-checkpoint_dir = snapshot_download(
-    repo_id=REPO_ID,
-    repo_type="model",
-    allow_patterns=["checkpoint/*"]
-)
-checkpoint_path = os.path.join(checkpoint_dir, "checkpoint")
-print(f"Downloaded to: {checkpoint_path}")
-
-print("Patching sharding metadata (TPU -> CPU)...")
-# 创建临时副本来修改 sharding 信息
-patched_checkpoint_path = "/tmp/mini_gpt_checkpoint_patched"
-if os.path.exists(patched_checkpoint_path):
-    shutil.rmtree(patched_checkpoint_path)
-shutil.copytree(checkpoint_path, patched_checkpoint_path, dirs_exist_ok=True)
-
-# 修改 _sharding 文件
-sharding_path = os.path.join(patched_checkpoint_path, "_sharding")
-if os.path.exists(sharding_path):
-    with open(sharding_path, 'r') as f:
-        sharding_data = json.load(f)
-    
-    cpu_device = jax.devices('cpu')[0]
-    cpu_device_str = str(cpu_device)
-    
-    new_sharding = {}
-    for key, value in sharding_data.items():
-        value_dict = json.loads(value)
-        value_dict['device_str'] = cpu_device_str
-        new_sharding[key] = json.dumps(value_dict)
-    
-    with open(sharding_path, 'w') as f:
-        json.dump(new_sharding, f)
-    
-    print(f"✓ Patched {len(new_sharding)} sharding entries to use CPU")
+model = MiniGPT(**CONFIG)
 
-print("Loading checkpoint...")
-checkpointer = ocp.PyTreeCheckpointer()
+print("Initializing random parameters...")
+rng = jax.random.PRNGKey(42)
+dummy_input = jnp.ones((1, CONFIG["max_len"]), dtype=jnp.int32)
+params = model.init(rng, dummy_input, training=False)['params']
 
-try:
-    state = checkpointer.restore(patched_checkpoint_path)
-    if 'params' in state:
-        params = state['params']
-        print("✓ Model loaded successfully!")
-    else:
-        raise ValueError("params not found in checkpoint")
-except Exception as e:
-    print(f"Error: {e}")
-    print("Using randomly initialized parameters")
-    rng = jax.random.PRNGKey(0)
-    dummy_input = jnp.ones((1, CONFIG["max_len"]), dtype=jnp.int32)
-    params = model.init(rng, dummy_input, training=False)['params']
-
-print(f"Params: {sum(x.size for x in jax.tree.leaves(params)):,} parameters")
+print(f"✓ Model ready with {sum(x.size for x in jax.tree.leaves(params)):,} parameters")
+print("⚠️ Note: Using random parameters (trained weights pending checkpoint fix)")
 
 
 # ============================================================================
@@ -240,9 +178,11 @@ def gradio_generate(prompt, max_tokens, temperature):
 
 with gr.Blocks(title="Mini-GPT 文本生成", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # 🤖 Mini-GPT 文本生成
+    # 🤖 Mini-GPT 文本生成 (Demo)
+    
+    使用 JAX/Flax 构建的小型 GPT 模型。
     
-    使用 JAX/Flax 在 Kaggle TPU 上训练的小型 GPT 模型。支持中英文输入。
+    ⚠️ **当前使用随机初始化参数** - 训练好的模型 checkpoint 正在修复中。
     """)
     
     with gr.Row():

requirements.txt

@@ -2,7 +2,6 @@ gradio==4.44.0
 jax==0.4.35
 jaxlib==0.4.35
 flax==0.10.2
-orbax-checkpoint==0.10.2
 transformers==4.47.0
 huggingface_hub>=0.23.0
 numpy==1.26.4