Fix startcode issue

Signed-off-by: yuanwu <yuan.wu@intel.com>
This commit is contained in:
yuanwu 2024-11-26 08:55:42 +00:00
parent c345c734a7
commit 636cdb4c43
2 changed files with 4 additions and 4 deletions

View File

@ -174,12 +174,12 @@ def get_model(
model_type = config_dict["model_type"] model_type = config_dict["model_type"]
if model_type == "gpt_bigcode": if model_type == "gpt_bigcode":
return StarCoder(model_id, revision, dtype) return StarCoder(model_id=model_id, revision=revision, dtype=dtype)
if model_type == "bloom": if model_type == "bloom":
return BLOOM( return BLOOM(
model_id, model_id=model_id,
revision, revision=revision,
speculator=speculator, speculator=speculator,
dtype=dtype, dtype=dtype,
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,

View File

@ -623,7 +623,7 @@ class CausalLM(Model):
def __init__( def __init__(
self, self,
model_id: str, model_id: str,
model_class, model_class: Optional[Type[torch.nn.Module]] = None,
revision: Optional[str] = None, revision: Optional[str] = None,
quantize: Optional[str] = None, quantize: Optional[str] = None,
speculator: Optional[str] = None, speculator: Optional[str] = None,