29 lines
1.0 KiB
Python
29 lines
1.0 KiB
Python
"""Print greedy answers to fixed prompts from an OpenAI-compatible server, as JSON.
|
|
|
|
usage: python3 scripts/greedy.py http://127.0.0.1:8000 MODEL > answers.json
|
|
"""
|
|
|
|
import json
|
|
import sys
|
|
import urllib.request
|
|
|
|
PROMPTS = [
|
|
"Explain what a mixture-of-experts language model is in two sentences.",
|
|
"Write a Python function that returns the n-th Fibonacci number.",
|
|
"Name three prime numbers greater than 100.",
|
|
]
|
|
|
|
|
|
def complete(base_url, model, prompt):
|
|
body = json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0, "max_tokens": 64}).encode()
|
|
request = urllib.request.Request(f"{base_url}/v1/chat/completions", body,
|
|
{"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(request, timeout=1800) as response:
|
|
return json.load(response)["choices"][0]["message"]["content"]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
base_url, model = sys.argv[1], sys.argv[2]
|
|
print(json.dumps([complete(base_url, model, p) for p in PROMPTS], indent=1))
|