Update app.py

app.py

@@ -102,30 +102,27 @@ DEMO_DATA = [
 # ==========================================
 # PART 1: 剧本生成模型 (ScriptAgent)
 # ==========================================
-from swift.llm import get_model_tokenizer, get_template, inference
+from swift.llm import PtEngine, RequestConfig, InferRequest
+from swift.plugin import InferStats
 import torch
 
 # 全局变量
 MODEL_NAME = "XD-MU/ScriptAgent"
 LOCAL_MODEL_PATH = "./downloaded_models/ScriptAgent"
-OFFLOAD_FOLDER = "./offload"
-model = None  # 模型对象
-tokenizer = None  # 分词器对象
-template = None  # 模板对象
+engine = None  # InferEngine 对象
 
 # 确保目录存在
 os.makedirs(LOCAL_MODEL_PATH, exist_ok=True)
-os.makedirs(OFFLOAD_FOLDER, exist_ok=True)
 
 def load_llm_model():
-    """使用 SWIFT 加载 ScriptAgent 模型 - CPU优化版本"""
-    global model, tokenizer, template
-    if model is not None:
+    """使用 ms-swift 的 PtEngine 加载 ScriptAgent 模型"""
+    global engine
+    if engine is not None:
         return
     
     try:
         # 1. 检查本地是否已下载模型
-        if not os.path.exists(LOCAL_MODEL_PATH):
+        if not os.path.exists(os.path.join(LOCAL_MODEL_PATH, "config.json")):
             print(f"正在从 HuggingFace 下载模型到 {LOCAL_MODEL_PATH}...")
             snapshot_download(
                 repo_id=MODEL_NAME,
@@ -137,28 +134,21 @@ def load_llm_model():
         else:
             print(f"✅ 模型已存在: {LOCAL_MODEL_PATH}")
         
-        # 2. 使用 SWIFT 正确加载模型
-        print("正在使用 SWIFT 加载模型（CPU + 半精度优化）...")
+        # 2. 使用 ms-swift 的 PtEngine 加载模型
+        print("正在使用 ms-swift PtEngine 加载模型...")
         
-        # 🔥 关键修改：使用 get_model_tokenizer
-        model, tokenizer = get_model_tokenizer(
+        # 🔥 关键修改：使用 PtEngine
+        engine = PtEngine(
             model_id_or_path=LOCAL_MODEL_PATH,
-            torch_dtype=torch.float16,           # 半精度
+            torch_dtype=torch.float16,      # 半精度
+            max_batch_size=1,                # 批处理大小
+            device_map='cpu',                # CPU设备
             model_kwargs={
-                'device_map': 'cpu',              # CPU设备
-                'low_cpu_mem_usage': True,        # 低内存模式
-                'offload_folder': OFFLOAD_FOLDER, # 内存溢出卸载到磁盘
-            },
-            max_model_len=4096,                   # 限制上下文长度
+                'low_cpu_mem_usage': True,   # 低内存模式
+            }
         )
         
-        # 设置为评估模式
-        model.eval()
-        
-        # 获取模板
-        template = get_template(tokenizer=tokenizer, model=model)
-        
-        print("✅ SWIFT 模型加载完成（已启用内存优化）")
+        print("✅ ms-swift PtEngine 加载完成")
         
     except Exception as e:
         print(f"❌ 模型加载失败: {e}")
@@ -166,12 +156,12 @@ def load_llm_model():
         traceback.print_exc()
 
 def chat_with_scriptagent(user_input: str):
-    """使用 SWIFT 与 ScriptAgent 对话生成剧本"""
-    global model, tokenizer, template
+    """使用 ms-swift InferEngine 与 ScriptAgent 对话生成剧本"""
+    global engine
     
-    if model is None:
+    if engine is None:
         load_llm_model()
-        if model is None:
+        if engine is None:
             return "❌ 模型加载失败，请检查后台日志。"
     
     user_input = user_input.strip()
@@ -179,23 +169,33 @@ def chat_with_scriptagent(user_input: str):
         return "请输入内容"
 
     try:
-        print("🤖 正在使用 SWIFT 推理剧本...")
+        print("🤖 正在使用 ms-swift InferEngine 推理剧本...")
         
-        # 🔥 使用 SWIFT 的 inference 函数
-        response, _ = inference(
-            model=model,
-            tokenizer=tokenizer,
-            template=template,
-            query=user_input,
-            max_new_tokens=4096,        # 从8192降低到4096
-            temperature=0.7,
-            top_p=0.9,
-            repetition_penalty=1.1,
-            do_sample=True,
-            num_beams=1,                # 贪婪解码
+        # 🔥 使用 ms-swift 的推理方式
+        # 1. 构建消息格式
+        messages = [{'role': 'user', 'content': user_input}]
+        infer_request = InferRequest(messages=messages)
+        
+        # 2. 配置请求参数
+        request_config = RequestConfig(
+            max_tokens=4096,           # 最大生成token数
+            temperature=0.7,           # 温度参数
+            top_p=0.9,                 # top_p 采样
+            repetition_penalty=1.1,    # 重复惩罚
+            stream=False,              # 不使用流式输出
         )
         
+        # 3. 执行推理
+        metric = InferStats()
+        resp_list = engine.infer([infer_request], request_config, metrics=[metric])
+        
+        # 4. 提取结果
+        response = resp_list[0].choices[0].message.content
+        
+        # 5. 打印性能指标（可选）
+        print(f"✅ 生成完成 | 指标: {metric.compute()}")
         print(f"✅ 生成结果长度: {len(response)} 字符")
+        
         return response if response else "⚠️ 生成为空，请重试"
     
     except Exception as e: