ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

chore: 简化格式并更新文档

Browse Source
This commit is contained in:
ViperEkura
2026-03-31 00:28:58 +08:00
parent eb57e55fca
commit 50488bd659
14 changed files with 506 additions and 582 deletions
Download Patch File Download Diff File Expand all files Collapse all files
demo/download.py
+4 -5
View File
@@ -1,13 +1,12 @@
import os
from pathlib import Path
from huggingface_hub import snapshot_download
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROJECT_ROOT = Path(__file__).parent.parent
PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
if __name__ == "__main__":
snapshot_download(
repo_id="ViperEk/KHAOSZ",
local_dir=os.path.join(PROJECT_ROOT, "params"),
local_dir=PARAMETER_ROOT,
force_download=True,
)
demo/generate_ar.py
+7 -8
View File
@@ -1,20 +1,19 @@
import os
import torch
from pathlib import Path
from khaosz.config.param_config import ModelParameter
from khaosz.inference.core import disable_random_init
from khaosz.inference.generator import LoopGenerator, GenerationRequest
from khaosz.inference.generator import GeneratorFactory, GenerationRequest
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROJECT_ROOT = Path(__file__).parent.parent
PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
def generate_text():
with disable_random_init():
model_dir = os.path.join(PROJECT_ROOT, "params")
param = ModelParameter.load(model_dir)
param = ModelParameter.load(PARAMETER_ROOT)
param.to(device="cuda", dtype=torch.bfloat16)
param.to(device="cuda", dtype=torch.bfloat16)
query = input(">> ")
request = GenerationRequest(
@@ -26,7 +25,7 @@ def generate_text():
history=None,
system_prompt=None,
)
generator = LoopGenerator(param)
generator = GeneratorFactory.create(param, request)
response = generator.generate(request)
print(response)
demo/generate_batch.py
+9 -8
View File
@@ -1,19 +1,19 @@
import os
import torch
from pathlib import Path
from khaosz.config.param_config import ModelParameter
from khaosz.inference.core import disable_random_init
from khaosz.inference.generator import BatchGenerator, GenerationRequest
from khaosz.inference.generator import GeneratorFactory, GenerationRequest
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROJECT_ROOT = Path(__file__).parent.parent
PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
def batch_generate():
with disable_random_init():
model_dir = os.path.join(PROJECT_ROOT, "params")
param = ModelParameter.load(model_dir)
param.to(device="cuda", dtype=torch.bfloat16)
generator = BatchGenerator(param)
with disable_random_init():
param = ModelParameter.load(PARAMETER_ROOT)
param.to(device="cuda", dtype=torch.bfloat16)
inputs = [
"你好",
"请问什么是人工智能",
@@ -31,6 +31,7 @@ def batch_generate():
history=None,
system_prompt=None,
)
generator = GeneratorFactory.create(param, request)
responses = generator.generate(request)
for q, r in zip(inputs, responses):
demo/stream_chat.py
+7 -9
View File
@@ -1,21 +1,18 @@
import os
import torch
from pathlib import Path
from khaosz.config.param_config import ModelParameter
from khaosz.inference.core import disable_random_init
from khaosz.inference.generator import StreamGenerator, GenerationRequest
from khaosz.inference.generator import GeneratorFactory, GenerationRequest
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROJECT_ROOT = Path(__file__).parent.parent
PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
def chat():
with disable_random_init():
model_dir = os.path.join(PROJECT_ROOT, "params")
param = ModelParameter.load(model_dir)
param.to(device="cuda", dtype=torch.bfloat16)
generator = StreamGenerator(param)
param = ModelParameter.load(PARAMETER_ROOT)
param.to(device="cuda", dtype=torch.bfloat16)
history = []
while True:
@@ -32,6 +29,7 @@ def chat():
history=history,
system_prompt=None,
)
generator = GeneratorFactory.create(param, request)
response_size = 0
full_response = ""