Update app.py

app.py

@@ -1,36 +1,47 @@
-from flask import Flask, request, Response, render_template
+from fastapi import FastAPI, Request
+from fastapi.responses import HTMLResponse, StreamingResponse
+from fastapi.templating import Jinja2Templates
 from pmbl import PMBL
-import os
-import json
 
-app = Flask(__name__, template_folder='.')
+
+
+
+app = FastAPI()
 pmbl = PMBL("./PMB-7b.Q6_K.gguf")  # Replace with the path to your model
 
-@app.route("/chat", methods=['POST'])
-def chat():
+templates = Jinja2Templates(directory=".")
+
+@app.post("/chat")
+async def chat(request: Request):
     try:
-        data = request.get_json()
+        data = await request.json()
         user_input = data["user_input"]
         mode = data["mode"]
         history = pmbl.get_chat_history(mode, user_input)
         response_generator = pmbl.generate_response(user_input, history, mode)
-        return Response(response_generator, mimetype='text/plain')
+        return StreamingResponse(response_generator, media_type="text/plain")
     except Exception as e:
         print(f"[SYSTEM] Error: {str(e)}")
-        return json.dumps({"error": str(e)})
+        return {"error": str(e)}
 
-@app.route("/", methods=['GET'])
-def root():
-    return render_template("index.html")
+@app.get("/", response_class=HTMLResponse)
+async def root(request: Request):
+    return templates.TemplateResponse("index.html", {"request": request})
 
-@app.route("/sleep", methods=['POST'])
-def sleep():
+@app.post("/sleep")
+async def sleep():
     try:
         pmbl.sleep_mode()
-        return json.dumps({"message": "Sleep mode completed successfully"})
+        return {"message": "Sleep mode completed successfully"}
     except Exception as e:
         print(f"[SYSTEM] Error: {str(e)}")
-        return json.dumps({"error": str(e)})
+        return {"error": str(e)}
+
+# Remove the main function as Hugging Face Spaces will automatically run the app. Using here to test
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 8000)))
+    import uvicorn
+    import asyncio
+
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(uvicorn.run(app, host="0.0.0.0", port=1771))