Update app.py

app.py

@@ -1,11 +1,7 @@
-from flask import Flask, request, jsonify, Response
-from flask_cors import CORS
+import gradio as gr
 from llamafactory.chat import ChatModel
 from llamafactory.extras.misc import torch_gc
 
-app = Flask(__name__)
-CORS(app)
-
 args = dict(
     model_name_or_path="StevenChen16/llama3-8b-Lawyer",
     template="llama3",
@@ -25,20 +21,24 @@ 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.
 """
 
-@app.route('/query', methods=['POST'])
-def query_model():
-    user_input = request.json.get('prompt', '')
+def query_model(user_input):
     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')
+    
+    response = ""
+    for new_text in chat_model.stream_chat(messages):
+        response += new_text
+    
+    return response
+
+# Gradio interface
+interface = gr.Interface(
+    fn=query_model,
+    inputs="text",
+    outputs="text",
+    title="Legal AI Assistant",
+    description="Ask me any legal question related to US and Canada law.",
+)
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000)
+    interface.launch()