64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
# Load environment variables first
|
|
from dotenv import load_dotenv
|
|
|
|
from llm import Llm
|
|
|
|
load_dotenv()
|
|
|
|
import os
|
|
from typing import Any, Coroutine
|
|
import asyncio
|
|
|
|
from evals.config import EVALS_DIR
|
|
from evals.core import generate_code_for_image
|
|
from evals.utils import image_to_data_url
|
|
|
|
STACK = "html_tailwind"
|
|
# MODEL = Llm.CLAUDE_3_5_SONNET_2024_06_20
|
|
N = 2 # Number of outputs to generate
|
|
|
|
|
|
async def main():
|
|
INPUT_DIR = EVALS_DIR + "/inputs"
|
|
OUTPUT_DIR = EVALS_DIR + "/outputs"
|
|
|
|
# Get all the files in the directory (only grab pngs)
|
|
evals = [f for f in os.listdir(INPUT_DIR) if f.endswith(".png")]
|
|
|
|
tasks: list[Coroutine[Any, Any, str]] = []
|
|
for filename in evals:
|
|
filepath = os.path.join(INPUT_DIR, filename)
|
|
data_url = await image_to_data_url(filepath)
|
|
for n in range(N): # Generate N tasks for each input
|
|
if n == 0:
|
|
task = generate_code_for_image(
|
|
image_url=data_url,
|
|
stack=STACK,
|
|
model=Llm.CLAUDE_3_5_SONNET_2024_06_20,
|
|
)
|
|
else:
|
|
task = generate_code_for_image(
|
|
image_url=data_url, stack=STACK, model=Llm.GPT_4O_2024_05_13
|
|
)
|
|
tasks.append(task)
|
|
|
|
print(f"Generating {len(tasks)} codes")
|
|
|
|
results = await asyncio.gather(*tasks)
|
|
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
|
for i, content in enumerate(results):
|
|
# Calculate index for filename and output number
|
|
eval_index = i // N
|
|
output_number = i % N
|
|
filename = evals[eval_index]
|
|
# File name is derived from the original filename in evals with an added output number
|
|
output_filename = f"{os.path.splitext(filename)[0]}_{output_number}.html"
|
|
output_filepath = os.path.join(OUTPUT_DIR, output_filename)
|
|
with open(output_filepath, "w") as file:
|
|
file.write(content)
|
|
|
|
|
|
asyncio.run(main())
|