Update app.py

app.py

@@ -1,134 +1,65 @@
 import gradio as gr
 from transformers import AutoTokenizer, LlamaForCausalLM
 import torch
-import psutil
-import gc
-from typing import List, Tuple
-import logging
 
-# 配置日志
-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()
-
-# 模型配置
+# 使用 UrbanGPT 模型
 model_name = "bjdwh/UrbanGPT"
 
-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
+# 加载模型和分词器
+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
+)
 
 def generate_response(
-    message: str,
-    history: List[Tuple[str, str]],
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
+    message,
+    history: list[tuple[str, str]],
+    max_tokens,
+    temperature,
+    top_p,
 ):
-    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  # 添加最大长度限制
+    # 格式化输入
+    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
         )
-        
-        # 将输入移到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)}"
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # 如果有历史对话，需要提取最后的回复
+    if history:
+        response = response.split("Assistant: ")[-1].strip()
+    
+    yield response
 
 # 创建 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 聊天助手",
@@ -136,12 +67,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    # 启动前进行内存优化
-    optimize_memory()
-    # 添加自定义配置
-    demo.launch(
-        share=False,
-        debug=True,
-        server_name="0.0.0.0",
-        server_port=7860
-    )
+    demo.launch()