small fix

2026-09-09 11:13:26 +08:00 · 2019-04-11 20:29:30 -06:00
parent e6135b3a0a
commit 8424ffdfbe
1 changed files with 1 additions and 1 deletions
@@ -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