From 8424ffdfbe9385dd01b6a59ff43173ff15cb0700 Mon Sep 17 00:00:00 2001 From: Nick Walton Date: Thu, 11 Apr 2019 20:29:30 -0600 Subject: [PATCH] small fix --- main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/main.py b/main.py index 5c0f991..041c0cc 100644 --- a/main.py +++ b/main.py @@ -68,7 +68,7 @@ def predict(context_tokens): def generate(prompt): context_tokens = enc.encode(prompt) pred = predict(context_tokens) - pred = pred[0][len(context_tokens):] + pred = pred[0]["output"][len(context_tokens):] output = enc.decode(pred) return output