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
      )