prain: skeleton

scripts/greedy.py

@@ -0,0 +1,28 @@
+"""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))