reafactor: 修改ModelParameter

2026-03-31 16:00:55 +08:00
parent 80c0b20877
commit 9f1561afe7
9 changed files with 48 additions and 60 deletions
@@ -4,23 +4,19 @@ import argparse

 from astrai.config.param_config import ModelParameter
 from astrai.inference.generator import BatchGenerator, GenerationRequest
-from astrai.inference.core import disable_random_init


 def processor(
    model_dir: str,
    input_json_file: str,
    output_json_file: str,
-    batch_size: int,
    temperature: float,
    top_k: int,
    top_p: float,
    question_key: str,
    response_key: str,
 ):
-    with disable_random_init():
-        param = ModelParameter.load(model_dir)
-
+    param = ModelParameter.load(model_dir, disable_init=True)
    param.to(device="cuda", dtype=torch.bfloat16)
    generator = BatchGenerator(param)

@@ -7,7 +7,6 @@ import tqdm

 from torch import Tensor
 from astrai.config.param_config import ModelParameter
-from astrai.inference.core import disable_random_init


 def compute_perplexity(
@@ -42,9 +41,7 @@ def compute_perplexity(
 def process_file(
    model_dir: str, input_file: str, output_file: str, batch_size: int, text_key: str
 ):
-    with disable_random_init():
-        param = ModelParameter.load(model_dir)
-
+    param = ModelParameter.load(model_dir, disable_init=True)
    param.to(device="cuda", dtype=torch.bfloat16)
    model = param.model
    tokenizer = param.tokenizer