Update app.py

app.py

@@ -1,15 +1,42 @@
-import threading
-import subprocess
 import gradio as gr
-from flask import Flask, request, stream_with_context, Response
-from flask_cors import CORS
+import os
+from threading import Thread
 from llamafactory.chat import ChatModel
+from llamafactory.extras.misc import torch_gc
 
-# 定义 Flask 后端 (app2.py 的内容)
-flask_app = Flask(__name__)
-CORS(flask_app)
 
-# 设置模型参数
+DESCRIPTION = '''
+<div>
+<h1 style="text-align: center;">AI Lawyer</h1>
+</div>
+'''
+
+LICENSE = """
+<p/>
+---
+Built with model "StevenChen16/Llama3-8B-Lawyer", based on "meta-llama/Meta-Llama-3-8B"
+"""
+
+PLACEHOLDER = """
+<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
+   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">AI Lawyer</h1>
+   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Ask me anything about US and Canada law...</p>
+</div>
+"""
+
+css = """
+h1 {
+  text-align: center;
+  display: block;
+}
+#duplicate-button {
+  margin: auto;
+  color: white;
+  background: #1565c0;
+  border-radius: 100vh;
+}
+"""
+
 args = dict(
     model_name_or_path="StevenChen16/llama3-8b-Lawyer",
     template="llama3",
@@ -17,11 +44,8 @@ args = dict(
     quantization_bit=8,
     use_unsloth=True,
 )
-
-# 加载模型
 chat_model = ChatModel(args)
 
-# 设置背景提示词
 background_prompt = """
 As an AI legal assistant, you are a highly trained expert in U.S. and Canadian law. Your purpose is to provide accurate, comprehensive, and professional legal information to assist users with a wide range of legal questions. When answering questions, you should actively ask questions to obtain more information, analyze from different perspectives, and explain your reasoning process to the user. Please adhere to the following guidelines:
 
@@ -70,52 +94,6 @@ Potential legal risks and coping strategies
 Please remember that your role is to provide general legal information and analysis, but also to actively guide and interact with the user during the conversation. If you find that necessary information is missing during the conversation, take the initiative to ask for it until you feel you have enough details to provide targeted analysis and advice. Now, please guide me step by step to describe the legal issues I am facing, according to the above requirements.
 """
 
-@flask_app.route('/predict', methods=['POST'])
-def predict():
-    data = request.get_json(force=True)
-    input_text = data.get("inputs", "")
-    prompt = f"{background_prompt}\n\nUser: {input_text}\nAssistant: "
-
-    def generate():
-        messages = [{"role": "user", "content": prompt}]
-        for new_text in chat_model.stream_chat(messages):
-            yield new_text + ' '  # 添加空格以区分单词
-
-    return Response(stream_with_context(generate()), content_type='text/plain')
-
-# 定义 Gradio 应用 (app.py 的内容)
-DESCRIPTION = '''
-<div>
-<h1 style="text-align: center;">AI Lawyer</h1>
-</div>
-'''
-
-LICENSE = """
-<p/>
----
-Built with model "StevenChen16/Llama3-8B-Lawyer", based on "meta-llama/Meta-Llama-3-8B"
-"""
-
-PLACEHOLDER = """
-<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
-   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">AI Lawyer</h1>
-   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Ask me anything about US and Canada law...</p>
-</div>
-"""
-
-css = """
-h1 {
-  text-align: center;
-  display: block;
-}
-#duplicate-button {
-  margin: auto;
-  color: white;
-  background: #1565c0;
-  border-radius: 100vh;
-}
-"""
-
 def query_model(user_input, history):
     combined_query = background_prompt + user_input
     messages = [{"role": "user", "content": combined_query}]
@@ -124,7 +102,7 @@ def query_model(user_input, history):
     for new_text in chat_model.stream_chat(messages, temperature=0.9):
         response += new_text
         yield response
-
+    
 # Gradio block
 chatbot = gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Gradio ChatInterface')
 
@@ -145,13 +123,5 @@ with gr.Blocks(css=css) as demo:
     )
     gr.Markdown(LICENSE)
 
-# 启动 Flask 应用的线程
-def run_flask():
-    flask_app.run(host='0.0.0.0', port=6006)
-
-thread = threading.Thread(target=run_flask)
-thread.start()
-
-# 启动 Gradio 应用
 if __name__ == "__main__":
     demo.launch(share=True)