116 lines
3.8 KiB
Python
116 lines
3.8 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, overwrite: bool = False) -> ExitCode:
|
|
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 not isinstance(previous, dict) or not isinstance(previous.get("jobs"), dict):
|
|
raise CliError(f"Invalid batch manifest: {workspace.batch_jobs_file}")
|
|
manifest["jobs"] = previous["jobs"]
|
|
started = 0
|
|
skipped = 0
|
|
for tier, file_path, model_id, display_name in batches:
|
|
existing = manifest["jobs"].get(tier)
|
|
if tier in manifest["jobs"]:
|
|
if (not isinstance(existing, dict)
|
|
or not isinstance(existing.get("name"), str)
|
|
or not existing["name"].strip()):
|
|
raise CliError(f"Invalid batch job in {workspace.batch_jobs_file}: {tier}")
|
|
if not overwrite:
|
|
print(f"Skipping {tier}: already submitted as {existing['name']}.")
|
|
skipped += 1
|
|
continue
|
|
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
|
|
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)
|
|
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 and not skipped:
|
|
print("No non-empty batch request files were found.")
|
|
return ExitCode.PARTIAL
|
|
return ExitCode.SUCCESS
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = evaluation_parser("Upload correction JSONL files and start Gemini batches")
|
|
parser.add_argument(
|
|
"--overwrite", action="store_true",
|
|
help="Submit new jobs even for recorded tiers, replacing their saved job IDs",
|
|
)
|
|
return parser
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
return execute(parser, argv, lambda args: run(
|
|
workspace_from_args(args), overwrite=args.overwrite,
|
|
))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|