108 lines
3.4 KiB
Python
108 lines
3.4 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)
|
|
|
|
|
|
# async def text_main():
|
|
# OUTPUT_DIR = EVALS_DIR + "/outputs"
|
|
|
|
# GENERAL_TEXT_V1 = [
|
|
# "Login form",
|
|
# "Simple notification",
|
|
# "button",
|
|
# "saas dashboard",
|
|
# "landing page for barber shop",
|
|
# ]
|
|
|
|
# tasks: list[Coroutine[Any, Any, str]] = []
|
|
# for prompt in GENERAL_TEXT_V1:
|
|
# for n in range(N): # Generate N tasks for each input
|
|
# if n == 0:
|
|
# task = generate_code_for_text(
|
|
# text=prompt,
|
|
# stack=STACK,
|
|
# model=Llm.CLAUDE_3_5_SONNET_2024_06_20,
|
|
# )
|
|
# else:
|
|
# task = generate_code_for_text(
|
|
# text=prompt, 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 = GENERAL_TEXT_V1[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())
|