Use OPTIMIZE_MODEL = GPT_4_128k()

This commit is contained in:
afik.cohen 2023-11-15 17:33:27 -08:00
parent 8a41cbaddb
commit 205e360e3e
4 changed files with 8 additions and 4 deletions

View file

@ -2,7 +2,7 @@ import os
from typing import Optional, Union
def get_codeflash_api_key() -> Union[str | None]:
def get_codeflash_api_key() -> Optional[str]:
api_key = os.environ.get("CODEFLASH_API_KEY")
if not api_key:
raise EnvironmentError(f"CODEFLASH_API_KEY not found in environment variables")

View file

@ -89,7 +89,9 @@ def parse_args():
if key in pyproject_config and getattr(args, key.replace("-", "_")) is None:
setattr(args, key.replace("-", "_"), pyproject_config[key])
assert os.path.isdir(args.root), "--root must be a valid directory"
assert os.path.isdir(args.test_root), "--test_root must be a valid directory"
assert os.path.isdir(
args.test_root
), f"--test-root must be a valid directory; {args.test_root} is not a directory"
args.root = os.path.realpath(args.root)
args.test_root = os.path.realpath(args.test_root)
if not hasattr(args, "all"):

View file

@ -1,7 +1,8 @@
import openai
import re
from typing import List, Tuple, Union
import openai
from codeflash.verification import OPTIMIZE_MODEL
def optimize_python_code(
@ -35,7 +36,7 @@ def optimize_python_code(
retries = 3
while retries > 0:
try:
output = openai.ChatCompletion.create(model="gpt-4", messages=messages, n=n)
output = openai.ChatCompletion.create(model=OPTIMIZE_MODEL.name, messages=messages, n=n)
retries -= 1
break
except Exception as e:

View file

@ -40,3 +40,4 @@ class GPT_3_5_Turbo(LLM):
EXPLAIN_MODEL: LLM = GPT_4_128k()
PLAN_MODEL = GPT_4_128k()
EXECUTE_MODEL = GPT_4_128k()
OPTIMIZE_MODEL = GPT_4_128k()