mk cov from api

codecritic/cli/mk_cov_from_api.py

+import argparse
+from itertools import chain
+import os
+from tqdm import tqdm
+from openai import OpenAI
+
+from codecritic.data.cov import (
+    convert_preference_to_vot_prompt,
+    convert_cov_to_cov_dataset,
+)
+from codecritic.data.utils import save_jsonl_dataset
+from codecritic.utils.json import load_json
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str)
+    parser.add_argument("--preference_dataset", type=str)
+    parser.add_argument("--output_dir", type=str)
+    args = parser.parse_args()
+
+    preference_dataset = load_json(args.preference_dataset)
+    cov_prompts = [convert_preference_to_vot_prompt(x) for x in preference_dataset]
+    cov_prompts = list(chain(*cov_prompts))
+
+
+    client = OpenAI(
+        base_url="https://api.deepseek.com/",
+        api_key=os.environ["DEEPSEEK_API_KEY"]
+    )
+
+    covs = []
+    for cov_prompt in tqdm(cov_prompts):
+        completion = client.chat.completions.create(
+            model="deepseek-chat",
+            messages=cov_prompt["messages"],
+            temperature=0,
+            max_tokens=2048
+        )
+        content = completion.choices[0].message.content
+        cov_prompt["messages"].append({"role": "assistant", "content": content})
+        covs.append(cov_prompt)
+
+    dataset = list(map(convert_cov_to_cov_dataset, covs))
+    save_jsonl_dataset(dataset, args.output_dir)
\ No newline at end of file