# TextSynth text completion example
import sys
import argparse
import requests
import json
import time

parser = argparse.ArgumentParser(description='TextSynth text completion example')
parser.add_argument('--url', type=str, default = "http://localhost:8080", help="server URL")
parser.add_argument('--api_key', type=str, default = "")
parser.add_argument('--model', type=str, default = "gpt2_117M", help = "model name")
parser.add_argument('-l', '--max_tokens', type=int, default = 100, help = "maximum number of output tokens")
parser.add_argument('-n', action='store_true', help = "don't output the prompt")
parser.add_argument('prompt', type=str, help="generation prompt")

args = parser.parse_args()

def make_request(path, json):
    response = requests.post(args.url + path, headers = { "Authorization": "Bearer " + args.api_key }, json = json)
    return response

def main():
    req = { "prompt": args.prompt, "max_tokens": args.max_tokens }
        
    result = make_request("/v1/engines/" + args.model + "/completions", req);
    if result.status_code != 200:
        print("Request error:", result.text)
        sys.exit(1)

    resp = result.json()
    if not args.n:
        print(args.prompt, end="");
    print(resp["text"])

main()