97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from collections.abc import Sequence
|
|
|
|
from google import genai
|
|
from google.genai import types
|
|
|
|
from copienator import configuration as config
|
|
from copienator import (
|
|
CliError,
|
|
EvaluationWorkspace,
|
|
ExitCode,
|
|
atomic_write_json,
|
|
evaluation_parser,
|
|
execute,
|
|
read_json,
|
|
workspace_from_args,
|
|
)
|
|
|
|
|
|
def run(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
|
|
if client is None:
|
|
if not config.API_KEY:
|
|
raise CliError("GEMINI_API_KEY is not configured")
|
|
client = genai.Client(api_key=config.API_KEY)
|
|
batches = (
|
|
(
|
|
"flash",
|
|
workspace.root / "batch_requests_flash.jsonl",
|
|
config.MODEL_FLASH_ID,
|
|
f"flash-correction-{workspace.name}",
|
|
),
|
|
(
|
|
"pro",
|
|
workspace.root / "batch_requests_pro.jsonl",
|
|
config.MODEL_PRO_ID,
|
|
f"pro-correction-{workspace.name}",
|
|
),
|
|
)
|
|
manifest = {
|
|
"version": 1,
|
|
"evaluation": workspace.name,
|
|
"jobs": {},
|
|
}
|
|
if workspace.batch_jobs_file.is_file():
|
|
previous = read_json(workspace.batch_jobs_file)
|
|
if isinstance(previous, dict) and isinstance(previous.get("jobs"), dict):
|
|
manifest["jobs"] = previous["jobs"]
|
|
started = 0
|
|
for tier, file_path, model_id, display_name in batches:
|
|
if not file_path.is_file():
|
|
print(f"Skipping {model_id}: {file_path.name} does not exist.")
|
|
continue
|
|
if file_path.stat().st_size == 0:
|
|
print(f"Skipping {model_id}: {file_path.name} is empty.")
|
|
continue
|
|
print(f"Uploading {file_path.name} for model {model_id}...")
|
|
uploaded = client.files.upload(
|
|
file=str(file_path),
|
|
config=types.UploadFileConfig(
|
|
display_name=f"{display_name}-input",
|
|
mime_type="jsonl",
|
|
),
|
|
)
|
|
job = client.batches.create(
|
|
model=model_id,
|
|
src=uploaded.name,
|
|
config={"display_name": display_name},
|
|
)
|
|
started += 1
|
|
manifest["jobs"][tier] = {
|
|
"name": job.name,
|
|
"display_name": display_name,
|
|
"model": model_id,
|
|
"request_file": file_path.name,
|
|
}
|
|
atomic_write_json(workspace.batch_jobs_file, manifest)
|
|
print(f"Started batch job: {job.name}")
|
|
if not started:
|
|
print("No non-empty batch request files were found.")
|
|
return ExitCode.PARTIAL
|
|
return ExitCode.SUCCESS
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
return evaluation_parser("Upload correction JSONL files and start Gemini batches")
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
return execute(parser, argv, lambda args: run(workspace_from_args(args)))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|