Update app.py

app.py

@@ -1,86 +1,78 @@
 import os
-import spaces
 import json
 import subprocess
+import gradio as gr
 from llama_cpp import Llama
 from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
 from llama_cpp_agent.providers import LlamaCppPythonProvider
 from llama_cpp_agent.chat_history import BasicChatHistory
 from llama_cpp_agent.chat_history.messages import Roles
-import gradio as gr
 from huggingface_hub import hf_hub_download
-import logging
-
-# 設置日誌
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
 
-# 設置參數
-REPO_ID = os.getenv("REPO_ID", "SakuraLLM/Sakura-14B-Qwen2beta-v0.9.2-GGUF")
-FILENAME = os.getenv("FILENAME", "sakura-14b-qwen2beta-v0.9.2-q4km.gguf")
-SYSTEM_MESSAGE = os.getenv("SYSTEM_MESSAGE", "你是一个轻小说翻译模型，可以流畅通顺地使用给定的术语表以日本轻小说的风格将日文翻译成简体中文，并联系上下文正确使用人称代词，注意不要混淆使役态和被动态的主语和宾语，不要擅自添加原文中没有的代词，也不要擅自增加或减少换行。")
+# 環境變量或配置文件管理
+REPO_ID = "SakuraLLM/Sakura-14B-Qwen2beta-v0.9.2-GGUF"
+FILENAME = "sakura-14b-qwen2beta-v0.9.2-q4km.gguf"
 MODEL_DIR = "./models"
+DEFAULT_SYSTEM_MESSAGE = ("你是一个轻小说翻译模型，可以流畅通顺地使用给定的术语表以日本轻小说的风格将日文翻译成简体中文，"
+                          "并联系上下文正确使用人称代词，注意不要混淆使役态和被动态的主语和宾语，不要擅自添加原文中没有的代词，"
+                          "也不要擅自增加或减少换行。")
 
-# 下載模型
-def download_model(repo_id, filename, local_dir):
-    logger.info(f"Downloading model {filename} from {repo_id} to {local_dir}")
-    hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)
+llm = None
+llm_model = None
 
-# 初始化 Llama 模型
-def initialize_llama(model_path):
-    logger.info(f"Initializing Llama model from {model_path}")
-    return Llama(
-        model_path=model_path,
-        flash_attn=True,
-        n_gpu_layers=81,
-        n_batch=1024,
-        n_ctx=8192,
-    )
+def download_model(repo_id, filename, local_dir):
+    """下載模型"""
+    try:
+        hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            local_dir=local_dir
+        )
+    except Exception as e:
+        print(f"下載模型失敗: {e}")
 
-# 處理歷史消息
-def process_history(history):
-    messages = BasicChatHistory()
-    for msn in history:
-        user = {
-            'role': Roles.user,
-            'content': "根据以下术语表（可以为空）：\n" + "将下面的日文文本根据上述术语表的对应关系和备注翻译成中文，并且列印出使用哪些术语表：" + msn[0]
-        }
-        assistant = {
-            'role': Roles.assistant,
-            'content': msn[1]
-        }
-        messages.add_message(user)
-        messages.add_message(assistant)
-    return messages
+def load_model(model_path, model):
+    """加載模型"""
+    global llm
+    global llm_model
+    
+    if llm is None or llm_model != model:
+        try:
+            llm = Llama(
+                model_path=model_path,
+                flash_attn=True,
+                n_gpu_layers=81,
+                n_batch=1024,
+                n_ctx=8192,
+            )
+            llm_model = model
+        except Exception as e:
+            print(f"加載模型失敗: {e}")
 
-# 主函數
-@spaces.GPU(duration=120)
 def respond(
     message,
-    history: list[tuple[str, str]],
+    history,
     model=FILENAME,
-    system_message=SYSTEM_MESSAGE,
+    system_message=DEFAULT_SYSTEM_MESSAGE,
     max_tokens=4096,
     temperature=0.1,
     top_p=0.3,
     top_k=40,
     repeat_penalty=1.1,
 ):
-    global llm
-    global llm_model
-
-    if llm is None or llm_model != model:
-        llm = initialize_llama(f"{MODEL_DIR}/{model}")
-        llm_model = model
-
+    """處理回應"""
+    chat_template = MessagesFormatterType.GEMMA_2
+    
+    load_model(f"{MODEL_DIR}/{model}", model)
+    
     provider = LlamaCppPythonProvider(llm)
     agent = LlamaCppAgent(
         provider,
-        system_prompt=system_message,
-        predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
+        system_prompt=f"{system_message}",
+        predefined_messages_formatter_type=chat_template,
         debug_output=True
     )
-
+    
     settings = provider.get_provider_default_settings()
     settings.temperature = temperature
     settings.top_k = top_k
@@ -88,9 +80,21 @@ def respond(
     settings.max_tokens = max_tokens
     settings.repeat_penalty = repeat_penalty
     settings.stream = True
-
-    messages = process_history(history)
-
+    
+    messages = BasicChatHistory()
+    
+    for msn in history:
+        user = {
+            'role': Roles.user,
+            'content': "根据以下术语表（可以为空）：\n" + "将下面的日文文本根据上述术语表���对应关系和备注翻译成中文，并且列印出使用哪些术语表：" + msn[0]
+        }
+        assistant = {
+            'role': Roles.assistant,
+            'content': msn[1]
+        }
+        messages.add_message(user)
+        messages.add_message(assistant)
+    
     stream = agent.get_chat_response(
         message,
         llm_sampling_settings=settings,
@@ -98,14 +102,14 @@ def respond(
         returns_streaming_generator=True,
         print_output=False
     )
-
+    
     outputs = ""
     for output in stream:
         outputs += output
     outputs = outputs.replace(system_message, '')
+    
     yield outputs
 
-# Gradio 接口
 description = """<p align="center">Defaults to Sakura-14B-Qwen2beta</p>
 <p><center>
 <a href="https://huggingface.co/SakuraLLM/Sakura-14B-Qwen2beta-v0.9.2-GGUF" target="_blank">[Sakura-14B-Qwen2beta Model]</a>
@@ -129,4 +133,4 @@ demo = gr.ChatInterface(
 
 if __name__ == "__main__":
     download_model(REPO_ID, FILENAME, MODEL_DIR)
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()