PMBeta / app.py
Sergidev's picture
Update app.py
3d08955 verified
raw
history blame
1.21 kB
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pmbl import PMBL
import os
app = FastAPI()
pmbl = PMBL("./PMB-7b.Q6_K.gguf") # Replace with the path to your model
app.mount("/", StaticFiles(directory="static"), name="static")
=
@app.post("/chat")
async def chat(request: Request):
try:
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 StreamingResponse(response_generator, media_type="text/plain")
except Exception as e:
print(f"[SYSTEM] Error: {str(e)}")
return {"error": str(e)}
@app.get("/", response_class=HTMLResponse)
async def root(request: Request):
return templates.TemplateResponse("index.html", {"request": request})
@app.post("/sleep")
async def sleep():
try:
pmbl.sleep_mode()
return {"message": "Sleep mode completed successfully"}
except Exception as e:
print(f"[SYSTEM] Error: {str(e)}")
return {"error": str(e)}