LIVE-Potassium
potassium.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
@app.init
def init():
model = torch.load('gptj.pt')
context = {
"model": model
}
return context
@app.handler()
def handler(context, request):
prompt = request.json.get("prompt")
model = context.get("model")
outputs = model(prompt)
return Response(
json = {"outputs": outputs},
status=200
)
