update app.py

app.py

@@ -1,17 +1,23 @@
-import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from flask import Flask, request, jsonify, Response
+from flask_cors import CORS
+from llamafactory.chat import ChatModel
+from llamafactory.extras.misc import torch_gc
 
-# 加载模型和tokenizer
-model_name = "StevenChen16/llama3-8b-Lawyer"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+app = Flask(__name__)
+CORS(app)
 
-# 检查是否有可用的GPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+args = dict(
+    model_name_or_path="StevenChen16/llama3-8b-Lawyer",
+    template="llama3",
+    finetuning_type="lora",
+    quantization_bit=8,
+    use_unsloth=True,
+)
+chat_model = ChatModel(args)
 
 background_prompt = """
 You are an advanced AI legal assistant trained to assist with a wide range of legal questions and issues. Your primary function is to provide accurate, comprehensive, and professional legal information based on U.S. and Canada law. Follow these guidelines when formulating responses:
+
 1. **Clarity and Precision**: Ensure that your responses are clear and precise. Use professional legal terminology, but explain complex legal concepts in a way that is understandable to individuals without a legal background.
 2. **Comprehensive Coverage**: Provide thorough answers that cover all relevant aspects of the question. Include explanations of legal principles, relevant statutes, case law, and their implications.
 3. **Contextual Relevance**: Tailor your responses to the specific context of the question asked. Provide examples or analogies where appropriate to illustrate legal concepts.
@@ -19,29 +25,20 @@ You are an advanced AI legal assistant trained to assist with a wide range of le
 5. **Professional Tone**: Maintain a professional and respectful tone in all responses. Ensure that your advice is legally sound and adheres to ethical standards.
 """
 
-def generate_response(input_text):
-    # 将输入文本与背景提示词结合
-    prompt = f"{background_prompt}\n\nUser: {input_text}\nAssistant: "
-    
-    # Tokenize 输入文本
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    
-    # 生成输出
-    outputs = model.generate(**inputs, max_new_tokens=300)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # 返回生成的响应
-    return response
-
-# 创建 Gradio 接口
-iface = gr.Interface(
-    fn=generate_response,
-    inputs="text",
-    outputs="text",
-    title="Legal Assistant",
-    description="Enter your legal question and get a response from the model.",
-)
+@app.route('/query', methods=['POST'])
+def query_model():
+    user_input = request.json.get('prompt', '')
+    combined_query = background_prompt + user_input
+    messages = [{"role": "user", "content": combined_query}]
+
+    def generate():
+        response = ""
+        for new_text in chat_model.stream_chat(messages):
+            response += new_text
+            yield f"data:{new_text}\n\n"
+        messages.append({"role": "assistant", "content": response})
+
+    return Response(generate(), mimetype='text/event-stream')
 
-# 运行应用
-if __name__ == "__main__":
-    iface.launch()
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000)