Update app.py

app.py

@@ -1,32 +1,65 @@
+import gradio as gr
+from huggingface_hub import hf_hub_download
 import os
 import subprocess
 
 # 模型下载链接
 model_url = "https://huggingface.co/CMLL/ZhongJing-2-1_8b-GGUF/resolve/main/ZhongJing1_5-1_8b-fp16.gguf"
 
-def install_packages():
-    subprocess.run(['apt-get', '-y', 'install', '-qq', 'aria2'], check=True)
-    subprocess.run(['git', 'clone', '-b', 'V20230828', 'https://github.com/Troyanovsky/text-generation-webui'], check=True)
-    os.chdir('/text-generation-webui')
-    subprocess.run(['pip', 'install', '-r', 'requirements.txt'], check=True)
-    subprocess.run(['pip', 'install', '-U', 'gradio==3.33.1'], check=True)
-    subprocess.run(['pip', 'uninstall', '-y', 'llama-cpp-python'], check=True)
-    os.environ['CMAKE_ARGS'] = "-DLLAMA_CUBLAS=on"
-    os.environ['FORCE_CMAKE'] = "1"
-    subprocess.run(['pip', 'install', 'llama-cpp-python', '--no-cache-dir'], check=True)
+# 下载模型函数
+def download_model(url, model_dir="models"):
+    os.makedirs(model_dir, exist_ok=True)
+    model_path = hf_hub_download(repo_id="CMLL/ZhongJing-2-1_8b-GGUF", filename="ZhongJing1_5-1_8b-fp16.gguf", local_dir=model_dir)
+    return model_path
 
-def download_model(model_url, model_name):
-    subprocess.run(['aria2c', '--console-log-level=error', '-c', '-x', '16', '-s', '16', '-k', '1M', model_url, '-d', '/text-generation-webui/models/', '-o', model_name], check=True)
+# 安装环境函数
+def setup_environment():
+    if not os.path.exists("llama.cpp"):
+        os.system("git clone https://github.com/ggerganov/llama.cpp.git")
+        os.system("cd llama.cpp && mkdir build && cd build && cmake .. && make")
 
-def run_server(model_name):
-    os.chdir('/text-generation-webui')
-    subprocess.run(['python', 'server.py', '--share', '--n-gpu-layers', '1000000000', '--model', model_name], check=True)
+    model_path = download_model(model_url)
+    
+    prompts_dir = "llama.cpp/prompts"
+    os.makedirs(prompts_dir, exist_ok=True)
+    with open(os.path.join(prompts_dir, "TcmChat.txt"), "w") as f:
+        f.write("You are a helpful TCM medical assistant named 仲景中医大语言模型.\n")
+    
+    return "Environment setup complete. Model downloaded and prompts file created."
+
+# 聊天函数
+def chat_with_model(user_input, history):
+    prompt = f"User: {user_input}\nAssistant:"
+    with open(os.path.join("llama.cpp/prompts", "TcmChat.txt"), "a") as f:
+        f.write(prompt + "\n")
+
+    # 执行命令并捕获输出
+    command = f"./llama.cpp/build/bin/main -m models/ZhongJing1_5-1_8b-fp16.gguf -n 256 --repeat_penalty 1.0 --color -i -r \"User:\" -f llama.cpp/prompts/TcmChat.txt"
+    result = subprocess.run(command, shell=True, capture_output=True, text=True)
+    
+    response = result.stdout.split("User:")[-1].strip()
+    history.append((user_input, response))
+    return history, history
+
+# 创建 Gradio 界面
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot()
+    state = gr.State([])
+
+    with gr.Row():
+        with gr.Column():
+            setup_btn = gr.Button("Setup Environment")
+            setup_output = gr.Textbox(label="Setup Output")
+        
+        with gr.Column():
+            user_input = gr.Textbox(show_label=False, placeholder="Enter your message...")
+            submit_btn = gr.Button("Submit")
+    
+    setup_btn.click(setup_environment, outputs=setup_output)
+    submit_btn.click(chat_with_model, [user_input, state], [chatbot, state])
 
 if __name__ == "__main__":
-    install_packages()
-    model_name = 'ZhongJing1_5-1_8b-fp16.gguf'
-    download_model(model_url, model_name)
-    run_server(model_name)
+    demo.launch()