115 lines
3.2 KiB
Python
115 lines
3.2 KiB
Python
from pathlib import Path
|
|
from typing import Any, TypedDict
|
|
|
|
import qai_hub.hub as hub
|
|
from qai_hub.client import CompileJob, Device, InferenceJob, Model, ProfileJob, QuantizeDtype, QuantizeJob
|
|
|
|
|
|
class ModelJobResult(TypedDict):
|
|
job: CompileJob | QuantizeJob
|
|
job_id: str
|
|
model: Model
|
|
model_id: str
|
|
|
|
|
|
class InferenceJobResult(TypedDict):
|
|
job: InferenceJob
|
|
job_id: str
|
|
outputs: Any
|
|
|
|
|
|
class ProfileJobResult(TypedDict):
|
|
job: ProfileJob
|
|
job_id: str
|
|
|
|
|
|
def _dataset_entries(inputs: dict[str, Any]) -> dict[str, list[Any]]:
|
|
return {name: value if isinstance(value, list) else [value] for name, value in inputs.items()}
|
|
|
|
|
|
def submit_compile_job(
|
|
model: Model,
|
|
device: Device,
|
|
input_specs: dict[str, tuple[tuple[int, ...], str]],
|
|
target_runtime: str,
|
|
options: str | None = None,
|
|
job_name: str | None = None,
|
|
) -> ModelJobResult:
|
|
compile_options = f"--target_runtime {target_runtime}"
|
|
if options:
|
|
compile_options = f"{compile_options} {options}"
|
|
|
|
job = hub.submit_compile_job(
|
|
model=model,
|
|
device=device,
|
|
name=job_name,
|
|
input_specs=input_specs,
|
|
options=compile_options,
|
|
)
|
|
target_model = job.get_target_model()
|
|
if target_model is None:
|
|
raise RuntimeError(f"Compile job {job.job_id} did not produce a target model.")
|
|
return {"job": job, "job_id": str(job.job_id), "model": target_model, "model_id": str(target_model.model_id)}
|
|
|
|
|
|
def submit_inference_job(
|
|
model: Model,
|
|
device: Device,
|
|
inputs: dict[str, Any],
|
|
output_dir: str | Path,
|
|
job_name: str | None = None,
|
|
) -> InferenceJobResult:
|
|
job = hub.submit_inference_job(
|
|
model=model,
|
|
device=device,
|
|
inputs=_dataset_entries(inputs),
|
|
name=job_name,
|
|
)
|
|
out = Path(output_dir)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
data = job.download_output_data(str(out))
|
|
return {"job": job, "job_id": str(job.job_id), "outputs": data}
|
|
|
|
|
|
def submit_profile_job(
|
|
model: Model,
|
|
device: Device,
|
|
options: str | None = None,
|
|
job_name: str | None = None,
|
|
) -> ProfileJobResult:
|
|
job = hub.submit_profile_job(
|
|
model=model,
|
|
device=device,
|
|
name=job_name,
|
|
options=options or "",
|
|
)
|
|
return {"job": job, "job_id": str(job.job_id)}
|
|
|
|
|
|
def submit_quantize_job(
|
|
model: Model,
|
|
calibration_data: dict[str, Any],
|
|
options: str | None = None,
|
|
job_name: str | None = None,
|
|
) -> ModelJobResult:
|
|
job = hub.submit_quantize_job(
|
|
model=model,
|
|
calibration_data=_dataset_entries(calibration_data),
|
|
weights_dtype=QuantizeDtype.INT8,
|
|
activations_dtype=QuantizeDtype.INT8,
|
|
name=job_name,
|
|
options=options or "",
|
|
)
|
|
target_model = job.get_target_model()
|
|
if target_model is None:
|
|
raise RuntimeError(f"Quantize job {job.job_id} did not produce a target model.")
|
|
return {"job": job, "job_id": str(job.job_id), "model": target_model, "model_id": str(target_model.model_id)}
|
|
|
|
|
|
def download_model(model_id: str, output_path: str | Path) -> str:
|
|
dest = Path(output_path)
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
model = hub.get_model(model_id)
|
|
result = model.download(str(dest))
|
|
return str(result or dest)
|