Update app.py

app.py

@@ -1,65 +1,134 @@
 import gradio as gr
 from transformers import AutoTokenizer, LlamaForCausalLM
 import torch
+import psutil
+import gc
+from typing import List, Tuple
+import logging
 
-# 使用 UrbanGPT 模型
+# 配置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def print_memory_usage():
+    """监控内存使用情况"""
+    process = psutil.Process()
+    cpu_mem = process.memory_info().rss / 1024 / 1024
+    gpu_mem = torch.cuda.memory_allocated() / 1024 / 1024 if torch.cuda.is_available() else 0
+    logger.info(f"CPU Memory: {cpu_mem:.2f}MB, GPU Memory: {gpu_mem:.2f}MB")
+
+def optimize_memory():
+    """优化内存使用"""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    print_memory_usage()
+
+# 模型配置
 model_name = "bjdwh/UrbanGPT"
 
-# 加载模型和分词器
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = LlamaForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True,
-    trust_remote_code=True
-)
+try:
+    # 加载模型和分词器
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        trust_remote_code=True
+    )
+    
+    # 使用 8bit 量化加载模型
+    model = LlamaForCausalLM.from_pretrained(
+        model_name,
+        load_in_8bit=True,          # 启用8bit量化
+        torch_dtype=torch.float16,   # 使用半精度
+        low_cpu_mem_usage=True,
+        trust_remote_code=True,
+        device_map="auto"           # 自动设备映射
+    )
+    
+    # 启用梯度检查点
+    model.gradient_checkpointing_enable()
+    
+except Exception as e:
+    logger.error(f"模型加载失败: {str(e)}")
+    raise
 
 def generate_response(
-    message,
-    history: list[tuple[str, str]],
-    max_tokens,
-    temperature,
-    top_p,
+    message: str,
+    history: List[Tuple[str, str]],
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
 ):
-    # 格式化输入
-    input_text = message
-    if history:
-        input_text = "\n".join([f"User: {h[0]}\nAssistant: {h[1]}" for h in history]) + f"\nUser: {message}"
-    
-    # 编码输入
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True)
-    
-    # 生成回复
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs["input_ids"],
-            max_length=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            num_return_sequences=1,
-            pad_token_id=tokenizer.eos_token_id
+    try:
+        optimize_memory()  # 优化内存使用
+        
+        # 格式化输入
+        input_text = message
+        if history:
+            input_text = "\n".join([f"User: {h[0]}\nAssistant: {h[1]}" for h in history]) + f"\nUser: {message}"
+        
+        # 编码输入
+        inputs = tokenizer(
+            input_text,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=2048  # 添加最大长度限制
         )
-    
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # 如果有历史对话，需要提取最后的回复
-    if history:
-        response = response.split("Assistant: ")[-1].strip()
-    
-    yield response
+        
+        # 将输入移到GPU（如果可用）
+        if torch.cuda.is_available():
+            inputs = {k: v.cuda() for k, v in inputs.items()}
+        
+        # 生成回复
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs["input_ids"],
+                max_length=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                num_return_sequences=1,
+                pad_token_id=tokenizer.eos_token_id,
+                do_sample=True,      # 启用采样
+                repetition_penalty=1.2  # 添加重复惩罚
+            )
+        
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # 提取最后的回复
+        if history:
+            response = response.split("Assistant: ")[-1].strip()
+        
+        optimize_memory()  # 生成后再次优化内存
+        yield response
+        
+    except Exception as e:
+        logger.error(f"生成回复时发生错误: {str(e)}")
+        yield f"抱歉，生成回复时发生错误: {str(e)}"
 
 # 创建 Gradio 界面
 demo = gr.ChatInterface(
     generate_response,
     additional_inputs=[
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="生成最大长度"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="温度"),
+        gr.Slider(
+            minimum=1,
+            maximum=2048,
+            value=512,
+            step=1,
+            label="生成最大长度"
+        ),
+        gr.Slider(
+            minimum=0.1,
+            maximum=4.0,
+            value=0.7,
+            step=0.1,
+            label="温度"
+        ),
         gr.Slider(
             minimum=0.1,
             maximum=1.0,
             value=0.95,
             step=0.05,
-            label="Top-p (核采样)",
+            label="Top-p (核采样)"
         ),
     ],
     title="UrbanGPT 聊天助手",
@@ -67,4 +136,12 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    # 启动前进行内存优化
+    optimize_memory()
+    # 添加自定义配置
+    demo.launch(
+        share=False,
+        debug=True,
+        server_name="0.0.0.0",
+        server_port=7860
+    )