2022-10-28 11:24:00 -06:00
|
|
|
from text_generation.models.model import Model
|
2022-11-03 09:07:54 -06:00
|
|
|
from text_generation.models.bloom import BLOOM, BLOOMSharded
|
2022-10-28 11:24:00 -06:00
|
|
|
|
2022-11-03 09:07:54 -06:00
|
|
|
__all__ = ["Model", "BLOOM", "BLOOMSharded"]
|
2022-10-28 11:24:00 -06:00
|
|
|
|
|
|
|
|
|
|
|
def get_model(model_name: str, sharded: bool, quantize: bool) -> Model:
|
|
|
|
if model_name.startswith("bigscience/bloom"):
|
|
|
|
if sharded:
|
|
|
|
return BLOOMSharded(model_name, quantize)
|
|
|
|
else:
|
|
|
|
if quantize:
|
|
|
|
raise ValueError("quantization is not supported for non-sharded BLOOM")
|
2022-11-03 09:07:54 -06:00
|
|
|
return BLOOM(model_name)
|
2022-10-28 11:24:00 -06:00
|
|
|
else:
|
2022-11-03 09:07:54 -06:00
|
|
|
raise ValueError(f"model {model_name} is not supported yet")
|