Update app.py

app.py

@@ -47,6 +47,7 @@ MAX_TOKENS = _int_env("MAX_TOKENS", 60)                      # 生成 token 上
 FEW_SHOT_EXAMPLES_COUNT = _int_env("FEW_SHOT", 0 if DEVICE == "cpu" else 1)
 ENABLE_INDEX = str(os.getenv("ENABLE_INDEX", "0" if DEVICE == "cpu" else "1")).lower() in {"1", "true", "yes", "y"}
 EMBED_BATCH = _int_env("EMBED_BATCH", 8 if DEVICE == "cpu" else 16)
+N_BATCH = _int_env("N_BATCH", 128 if DEVICE == "cpu" else 256)
 
 # 使用 /tmp 作為暫存目錄
 TEMP_DIR = "/tmp/text_to_sql_cache"
@@ -58,6 +59,7 @@ print(f"數據集: {DATASET_REPO_ID}")
 print(f"嵌入模型: {EMBED_MODEL_NAME}")
 print(f"設備: {DEVICE} (USE_GPU={USE_GPU}, N_GPU_LAYERS={N_GPU_LAYERS})")
 print(f"THREADS={THREADS}, CTX={CTX}, MAX_TOKENS={MAX_TOKENS}, FEW_SHOT={FEW_SHOT_EXAMPLES_COUNT}, ENABLE_INDEX={ENABLE_INDEX}, EMBED_BATCH={EMBED_BATCH}")
+print(f"N_BATCH={N_BATCH}")
 print(f"暫存目錄: {TEMP_DIR}")
 print("=" * 60)
 
@@ -203,7 +205,7 @@ class TextToSQLSystem:
                 model_path=model_path,
                 n_ctx=CTX,       # 上下文長度（CPU 默認更小）
                 n_threads=THREADS,  # 使用多執行緒
-                n_batch=256,     # 批處理大小
+                n_batch=N_BATCH,     # 批處理大小（可配置）
                 verbose=False,
                 n_gpu_layers=ngl,  # 可選 GPU 加速
                 use_mmap=True,   # 使用內存映射減少內存占用
@@ -254,6 +256,7 @@ class TextToSQLSystem:
             # 清理垃圾收集
             gc.collect()
 
+            start_ts = datetime.now()
             output = self.llm(
                 prompt,
                 max_tokens=MAX_TOKENS,  # 生成長度可配置
@@ -262,6 +265,8 @@ class TextToSQLSystem:
                 echo=False,
                 stop=["```", ";", "\n\n", "</s>"],
             )
+            elapsed = (datetime.now() - start_ts).total_seconds()
+            self._log(f"推論耗時: {elapsed:.2f}s", "DEBUG")
 
             self._log(f"模型原始輸出: {str(output)[:200]}...", "DEBUG")