wehub-resource-sync
46 lines
1.1 KiB
Python
46 lines
1.1 KiB
Python
from pathlib import Path
|
|
|
|
from fouroversix import ModelQuantizationConfig
|
|
|
|
from ...resources import (
|
|
FOUROVERSIX_CACHE_PATH,
|
|
app,
|
|
cache_volume,
|
|
get_image,
|
|
hf_secret,
|
|
)
|
|
from .evaluator import PTQEvaluator
|
|
|
|
hp_img = get_image()
|
|
|
|
with hp_img.imports():
|
|
from transformers import AutoModelForCausalLM
|
|
|
|
|
|
@app.cls(
|
|
image=hp_img,
|
|
gpu="B200",
|
|
secrets=[hf_secret],
|
|
timeout=24 * 60 * 60,
|
|
volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume},
|
|
)
|
|
class HighPrecisionEvaluator(PTQEvaluator):
|
|
"""Evaluate a model while keeping it in high precision."""
|
|
|
|
def quantize_model(
|
|
self,
|
|
model_name: str,
|
|
*,
|
|
device: str,
|
|
save_path: Path, # noqa: ARG002
|
|
quantization_config: ModelQuantizationConfig, # noqa: ARG002
|
|
trust_remote_code: bool = False,
|
|
) -> "AutoModelForCausalLM":
|
|
"""Return a model without any quantization."""
|
|
|
|
return AutoModelForCausalLM.from_pretrained(
|
|
model_name,
|
|
device_map=device,
|
|
trust_remote_code=trust_remote_code,
|
|
)
|