Create app.py

app.py

@@ -0,0 +1,43 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+import spaces
+
+# 你想部署的模型路径
+model_id = "dealignai/Gemma-4-31B-JANG_4M-CRACK"
+
+@spaces.GPU
+def chat(message, history):
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    # 使用 4-bit 量化加载以适配云端环境
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, 
+        device_map="auto", 
+        load_in_4bit=True,
+        torch_dtype=torch.bfloat16
+    )
+    
+    inputs = tokenizer(message, return_tensors="pt").to("cuda")
+    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    
+    generate_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        top_p=0.95,
+        temperature=0.7,
+    )
+    
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+    
+    partial_message = ""
+    for new_token in streamer:
+        partial_message += new_token
+        yield partial_message
+
+# 创建简单的聊天界面
+demo = gr.ChatInterface(fn=chat, title="Gemma 4 31B 数学逻辑测试")
+demo.launch()