reafactor: 修改ModelParameter

This commit is contained in:
2026-03-31 16:00:55 +08:00
parent 80c0b20877
commit 9f1561afe7
9 changed files with 48 additions and 60 deletions
+1 -1
View File
@@ -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,
)
+2 -5
View File
@@ -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(">> ")
+2 -5
View File
@@ -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 = [
"你好",
+2 -5
View File
@@ -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: