reafactor: 修改ModelParameter

scripts/demo/download.py

@@ -6,7 +6,7 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 if __name__ == "__main__":
     snapshot_download(
-        repo_id="ViperEk/AstrAI",
+        repo_id="ViperEk/KHAOSZ",
         local_dir=PARAMETER_ROOT,
         force_download=True,
     )

scripts/demo/generate_ar.py

@@ -1,7 +1,6 @@
 import torch
 from pathlib import Path
 from astrai.config.param_config import ModelParameter
-from astrai.inference.core import disable_random_init
 from astrai.inference.generator import GeneratorFactory, GenerationRequest
 
 PROJECT_ROOT = Path(__file__).parent.parent
@@ -9,10 +8,8 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
 def generate_text():
-
-    with disable_random_init():
-        param = ModelParameter.load(PARAMETER_ROOT)
-        param.to(device="cuda", dtype=torch.bfloat16)
+    param = ModelParameter.load(PARAMETER_ROOT, disable_init=True)
+    param.to(device="cuda", dtype=torch.bfloat16)
 
     query = input(">> ")
 

scripts/demo/generate_batch.py

@@ -1,7 +1,6 @@
 import torch
 from pathlib import Path
 from astrai.config.param_config import ModelParameter
-from astrai.inference.core import disable_random_init
 from astrai.inference.generator import GeneratorFactory, GenerationRequest
 
 PROJECT_ROOT = Path(__file__).parent.parent
@@ -9,10 +8,8 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
 def batch_generate():
-
-    with disable_random_init():
-        param = ModelParameter.load(PARAMETER_ROOT)
-        param.to(device="cuda", dtype=torch.bfloat16)
+    param = ModelParameter.load(PARAMETER_ROOT, disable_init=True)
+    param.to(device="cuda", dtype=torch.bfloat16)
 
     inputs = [
         "你好",

scripts/demo/stream_chat.py

@@ -1,7 +1,6 @@
 import torch
 from pathlib import Path
 from astrai.config.param_config import ModelParameter
-from astrai.inference.core import disable_random_init
 from astrai.inference.generator import GeneratorFactory, GenerationRequest
 
 PROJECT_ROOT = Path(__file__).parent.parent
@@ -9,10 +8,8 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
 def chat():
-
-    with disable_random_init():
-        param = ModelParameter.load(PARAMETER_ROOT)
-        param.to(device="cuda", dtype=torch.bfloat16)
+    param = ModelParameter.load(PARAMETER_ROOT, disable_init=True)
+    param.to(device="cuda", dtype=torch.bfloat16)
 
     history = []
     while True:

scripts/tools/generate.py

@@ -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)
 

scripts/tools/perplexity.py

@@ -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