From 2e009cf59a2f7ce2597d66155dcef7b4eb12afef Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Tue, 31 Mar 2026 09:34:11 +0800
Subject: [PATCH] =?UTF-8?q?chore:=20=E6=9B=B4=E6=96=B0=E9=A1=B9=E7=9B=AE?=
 =?UTF-8?q?=E5=90=8D=E7=A7=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 README.md                                    | 14 +++++++-------
 assets/docs/dataflow.md                      | 20 ++++++++++----------
 assets/docs/design.md                        |  2 +-
 assets/docs/params.md                        |  4 ++--
 {khaosz => astrai}/__init__.py               | 10 +++++-----
 {khaosz => astrai}/config/__init__.py        |  8 ++++----
 {khaosz => astrai}/config/model_config.py    |  0
 {khaosz => astrai}/config/param_config.py    |  6 +++---
 {khaosz => astrai}/config/schedule_config.py |  0
 {khaosz => astrai}/config/train_config.py    |  0
 {khaosz => astrai}/data/__init__.py          |  6 +++---
 {khaosz => astrai}/data/dataset.py           |  2 +-
 {khaosz => astrai}/data/sampler.py           |  0
 {khaosz => astrai}/data/serialization.py     |  2 +-
 {khaosz => astrai}/data/tokenizer.py         |  0
 {khaosz => astrai}/inference/__init__.py     |  4 ++--
 {khaosz => astrai}/inference/core.py         |  2 +-
 {khaosz => astrai}/inference/generator.py    |  4 ++--
 {khaosz => astrai}/model/__init__.py         |  4 ++--
 {khaosz => astrai}/model/module.py           |  0
 {khaosz => astrai}/model/transformer.py      |  4 ++--
 {khaosz => astrai}/parallel/__init__.py      |  4 ++--
 {khaosz => astrai}/parallel/module.py        |  0
 {khaosz => astrai}/parallel/setup.py         |  0
 {khaosz => astrai}/trainer/__init__.py       |  8 ++++----
 {khaosz => astrai}/trainer/metric_util.py    |  0
 {khaosz => astrai}/trainer/schedule.py       |  2 +-
 {khaosz => astrai}/trainer/strategy.py       |  0
 {khaosz => astrai}/trainer/train_callback.py |  8 ++++----
 {khaosz => astrai}/trainer/train_context.py  | 10 +++++-----
 {khaosz => astrai}/trainer/trainer.py        | 10 +++++-----
 demo/download.py                             |  2 +-
 demo/generate_ar.py                          |  6 +++---
 demo/generate_batch.py                       |  6 +++---
 demo/stream_chat.py                          |  6 +++---
 pyproject.toml                               |  6 +++---
 tests/conftest.py                            |  6 +++---
 tests/data/test_checkpoint.py                |  4 ++--
 tests/data/test_dataset.py                   |  4 ++--
 tests/data/test_sampler.py                   |  4 ++--
 tests/module/test_module.py                  | 10 +++++-----
 tests/module/test_tie_weight.py              |  4 ++--
 tests/test_parallel.py                       |  2 +-
 tests/trainer/test_callbacks.py              |  4 ++--
 tests/trainer/test_early_stopping.py         |  6 +++---
 tests/trainer/test_train_strategy.py         |  6 +++---
 tests/trainer/test_trainer.py                |  6 +++---
 tools/benchmark.py                           |  2 +-
 tools/generate.py                            |  6 +++---
 tools/perplexity.py                          |  4 ++--
 tools/train.py                               |  8 ++++----
 51 files changed, 118 insertions(+), 118 deletions(-)
 rename {khaosz => astrai}/__init__.py (69%)
 rename {khaosz => astrai}/config/__init__.py (63%)
 rename {khaosz => astrai}/config/model_config.py (100%)
 rename {khaosz => astrai}/config/param_config.py (94%)
 rename {khaosz => astrai}/config/schedule_config.py (100%)
 rename {khaosz => astrai}/config/train_config.py (100%)
 rename {khaosz => astrai}/data/__init__.py (79%)
 rename {khaosz => astrai}/data/dataset.py (99%)
 rename {khaosz => astrai}/data/sampler.py (100%)
 rename {khaosz => astrai}/data/serialization.py (98%)
 rename {khaosz => astrai}/data/tokenizer.py (100%)
 rename {khaosz => astrai}/inference/__init__.py (86%)
 rename {khaosz => astrai}/inference/core.py (99%)
 rename {khaosz => astrai}/inference/generator.py (98%)
 rename {khaosz => astrai}/model/__init__.py (62%)
 rename {khaosz => astrai}/model/module.py (100%)
 rename {khaosz => astrai}/model/transformer.py (98%)
 rename {khaosz => astrai}/parallel/__init__.py (76%)
 rename {khaosz => astrai}/parallel/module.py (100%)
 rename {khaosz => astrai}/parallel/setup.py (100%)
 rename {khaosz => astrai}/trainer/__init__.py (70%)
 rename {khaosz => astrai}/trainer/metric_util.py (100%)
 rename {khaosz => astrai}/trainer/schedule.py (99%)
 rename {khaosz => astrai}/trainer/strategy.py (100%)
 rename {khaosz => astrai}/trainer/train_callback.py (97%)
 rename {khaosz => astrai}/trainer/train_context.py (91%)
 rename {khaosz => astrai}/trainer/trainer.py (93%)

diff --git a/README.md b/README.md
index d8a9395..818d5c0 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
 <div align="center">
-  <img src="assets/images/project_logo.png" width="auto" alt="Logo">
+  <!-- <img src="assets/images/project_logo.png" width="auto" alt="Logo"> -->
   
-  <h1>KHAOSZ</h1>
+  <h1>AstrAI</h1>
   
   <div>
     <a href="#english">English</a> • 
@@ -48,8 +48,8 @@
 ### Installation
 
 ```bash
-git clone https://github.com/username/khaosz.git
-cd khaosz
+git clone https://github.com/ViperEkura/AstrAI.git
+cd AstrAI
 pip install -e .
 ```
 
@@ -95,8 +95,8 @@ python demo/generate_ar.py
 ### 安装
 
 ```bash
-git clone https://github.com/username/khaosz.git
-cd khaosz
+git clone https://github.com/ViperEkura/AstrAI.git
+cd AstrAI
 pip install -e .
 ```
 
@@ -143,7 +143,7 @@ python demo/generate_ar.py
 
 ### Download | 下载
 
-- [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ)
+- [HuggingFace](https://huggingface.co/ViperEk/AstrAI)
 - `python demo/download.py`
 
 ### Lincence | 许可证
diff --git a/assets/docs/dataflow.md b/assets/docs/dataflow.md
index 639b96b..f1e6bb7 100644
--- a/assets/docs/dataflow.md
+++ b/assets/docs/dataflow.md
@@ -1,16 +1,16 @@
-# KHAOSZ Data Flow Documentation
+# AstrAI Data Flow Documentation
 
-This document describes the data flow of the KHAOSZ project (a training and inference framework for autoregressive Transformer language models). It covers the complete flow from raw data to model training and inference.
+This document describes the data flow of the AstrAI project (a training and inference framework for autoregressive Transformer language models). It covers the complete flow from raw data to model training and inference.
 
 ## Overview
 
-KHAOSZ adopts a modular design with the following main components:
-- **Data Module** (`khaosz/data/`): Dataset, sampler, tokenizer, serialization tools
-- **Model Module** (`khaosz/model/`): Transformer model and its submodules
-- **Training Module** (`khaosz/trainer/`): Trainer, training context, strategies, schedulers
-- **Inference Module** (`khaosz/inference/`): Generation core, KV cache management, streaming generation
-- **Config Module** (`khaosz/config/`): Model, training, scheduler, and other configurations
-- **Parallel Module** (`khaosz/parallel/`): Distributed training support
+AstrAI adopts a modular design with the following main components:
+- **Data Module** (`astrai/data/`): Dataset, sampler, tokenizer, serialization tools
+- **Model Module** (`astrai/model/`): Transformer model and its submodules
+- **Training Module** (`astrai/trainer/`): Trainer, training context, strategies, schedulers
+- **Inference Module** (`astrai/inference/`): Generation core, KV cache management, streaming generation
+- **Config Module** (`astrai/config/`): Model, training, scheduler, and other configurations
+- **Parallel Module** (`astrai/parallel/`): Distributed training support
 
 The data flow can generally be divided into two main lines: **Training Data Flow** and **Inference Data Flow**.
 
@@ -199,7 +199,7 @@ flowchart LR
 
 ## Summary
 
-The data flow design of KHAOSZ reflects the characteristics of modularity, extensibility, and resumability. The training data flow supports large-scale distributed training through chunk loading, resumable sampling, gradient accumulation, and other mechanisms; the inference data flow achieves efficient text generation using KV cache and sampling strategies. Clear interfaces between modules facilitate customization and extension.
+The data flow design of AstrAI reflects the characteristics of modularity, extensibility, and resumability. The training data flow supports large-scale distributed training through chunk loading, resumable sampling, gradient accumulation, and other mechanisms; the inference data flow achieves efficient text generation using KV cache and sampling strategies. Clear interfaces between modules facilitate customization and extension.
 
 > Document Update Time: 2026-03-30  
 > Corresponding Code Version: Refer to version number defined in `pyproject.toml`
\ No newline at end of file
diff --git a/assets/docs/design.md b/assets/docs/design.md
index ba7bc61..e99328d 100644
--- a/assets/docs/design.md
+++ b/assets/docs/design.md
@@ -2,7 +2,7 @@
 
 There are many large language models on the market today, such as GPT, LLaMA, and others, with tens of billions or even hundreds of billions of parameters. But honestly, these models have extremely high hardware requirements, making them inaccessible for ordinary developers. I thought: **Can we create a model that is both useful and can run on ordinary computers?** This is also what most people currently hope for - a locally deployable AI project that achieves complete privatization while maintaining some level of intelligence.
 
-Thus, the KHAOSZ project was born - 1B parameters, Chinese-English bilingual, supporting dialogue, text generation, RAG retrieval, and the training code is open source!
+Thus, the AstrAI project was born - 1B parameters, Chinese-English bilingual, supporting dialogue, text generation, RAG retrieval, and the training code is open source!
 
 ## 2. System Architecture
 
diff --git a/assets/docs/params.md b/assets/docs/params.md
index 965812e..256bfa4 100644
--- a/assets/docs/params.md
+++ b/assets/docs/params.md
@@ -83,8 +83,8 @@
 ### Usage Example
 
 ```python
-from khaosz.config.param_config import ModelParameter
-from khaosz.inference.generator import StreamGenerator, GenerationRequest
+from astrai.config.param_config import ModelParameter
+from astrai.inference.generator import StreamGenerator, GenerationRequest
 
 # Load model
 param = ModelParameter.load("your_model_dir")
diff --git a/khaosz/__init__.py b/astrai/__init__.py
similarity index 69%
rename from khaosz/__init__.py
rename to astrai/__init__.py
index cb88bb8..05bc4f6 100644
--- a/khaosz/__init__.py
+++ b/astrai/__init__.py
@@ -1,13 +1,13 @@
 __version__ = "1.3.2"
 __author__ = "ViperEkura"
 
-from khaosz.config import (
+from astrai.config import (
     ModelConfig,
     TrainConfig,
 )
-from khaosz.model.transformer import Transformer
-from khaosz.data import DatasetLoader, BpeTokenizer
-from khaosz.inference.generator import (
+from astrai.model.transformer import Transformer
+from astrai.data import DatasetLoader, BpeTokenizer
+from astrai.inference.generator import (
     GenerationRequest,
     LoopGenerator,
     StreamGenerator,
@@ -15,7 +15,7 @@ from khaosz.inference.generator import (
     EmbeddingEncoder,
     GeneratorFactory,
 )
-from khaosz.trainer import Trainer, StrategyFactory, SchedulerFactory
+from astrai.trainer import Trainer, StrategyFactory, SchedulerFactory
 
 __all__ = [
     "Transformer",
diff --git a/khaosz/config/__init__.py b/astrai/config/__init__.py
similarity index 63%
rename from khaosz/config/__init__.py
rename to astrai/config/__init__.py
index eb2a81b..e2509b0 100644
--- a/khaosz/config/__init__.py
+++ b/astrai/config/__init__.py
@@ -1,12 +1,12 @@
-from khaosz.config.model_config import ModelConfig
-from khaosz.config.param_config import BaseModelIO, ModelParameter
-from khaosz.config.schedule_config import (
+from astrai.config.model_config import ModelConfig
+from astrai.config.param_config import BaseModelIO, ModelParameter
+from astrai.config.schedule_config import (
     ScheduleConfig,
     CosineScheduleConfig,
     SGDRScheduleConfig,
     ScheduleConfigFactory,
 )
-from khaosz.config.train_config import TrainConfig
+from astrai.config.train_config import TrainConfig
 
 
 __all__ = [
diff --git a/khaosz/config/model_config.py b/astrai/config/model_config.py
similarity index 100%
rename from khaosz/config/model_config.py
rename to astrai/config/model_config.py
diff --git a/khaosz/config/param_config.py b/astrai/config/param_config.py
similarity index 94%
rename from khaosz/config/param_config.py
rename to astrai/config/param_config.py
index 2f61098..1a3f79d 100644
--- a/khaosz/config/param_config.py
+++ b/astrai/config/param_config.py
@@ -5,9 +5,9 @@ from dataclasses import dataclass, field
 from typing import Optional, Self, Union
 from pathlib import Path
 
-from khaosz.data.tokenizer import BpeTokenizer
-from khaosz.config.model_config import ModelConfig
-from khaosz.model.transformer import Transformer
+from astrai.data.tokenizer import BpeTokenizer
+from astrai.config.model_config import ModelConfig
+from astrai.model.transformer import Transformer
 
 
 @dataclass
diff --git a/khaosz/config/schedule_config.py b/astrai/config/schedule_config.py
similarity index 100%
rename from khaosz/config/schedule_config.py
rename to astrai/config/schedule_config.py
diff --git a/khaosz/config/train_config.py b/astrai/config/train_config.py
similarity index 100%
rename from khaosz/config/train_config.py
rename to astrai/config/train_config.py
diff --git a/khaosz/data/__init__.py b/astrai/data/__init__.py
similarity index 79%
rename from khaosz/data/__init__.py
rename to astrai/data/__init__.py
index c924023..6a526cb 100644
--- a/khaosz/data/__init__.py
+++ b/astrai/data/__init__.py
@@ -1,4 +1,4 @@
-from khaosz.data.dataset import (
+from astrai.data.dataset import (
     BaseDataset,
     SEQDataset,
     DPODataset,
@@ -9,8 +9,8 @@ from khaosz.data.dataset import (
     DatasetFactory,
 )
 
-from khaosz.data.tokenizer import BpeTokenizer
-from khaosz.data.sampler import ResumableDistributedSampler
+from astrai.data.tokenizer import BpeTokenizer
+from astrai.data.sampler import ResumableDistributedSampler
 
 __all__ = [
     # Base classes
diff --git a/khaosz/data/dataset.py b/astrai/data/dataset.py
similarity index 99%
rename from khaosz/data/dataset.py
rename to astrai/data/dataset.py
index b696644..60e70d9 100644
--- a/khaosz/data/dataset.py
+++ b/astrai/data/dataset.py
@@ -6,7 +6,7 @@ import bisect
 from abc import ABC, abstractmethod
 from torch import Tensor
 from torch.utils.data import Dataset
-from khaosz.data.serialization import load_h5
+from astrai.data.serialization import load_h5
 from typing import List, Dict, Optional, Union
 
 
diff --git a/khaosz/data/sampler.py b/astrai/data/sampler.py
similarity index 100%
rename from khaosz/data/sampler.py
rename to astrai/data/sampler.py
diff --git a/khaosz/data/serialization.py b/astrai/data/serialization.py
similarity index 98%
rename from khaosz/data/serialization.py
rename to astrai/data/serialization.py
index af6c859..4f52687 100644
--- a/khaosz/data/serialization.py
+++ b/astrai/data/serialization.py
@@ -8,7 +8,7 @@ import torch.distributed as dist
 from pathlib import Path
 from torch import Tensor
 from typing import Any, Dict, List
-from khaosz.parallel.setup import get_rank
+from astrai.parallel.setup import get_rank
 
 
 def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
diff --git a/khaosz/data/tokenizer.py b/astrai/data/tokenizer.py
similarity index 100%
rename from khaosz/data/tokenizer.py
rename to astrai/data/tokenizer.py
diff --git a/khaosz/inference/__init__.py b/astrai/inference/__init__.py
similarity index 86%
rename from khaosz/inference/__init__.py
rename to astrai/inference/__init__.py
index 03549b5..28e2174 100644
--- a/khaosz/inference/__init__.py
+++ b/astrai/inference/__init__.py
@@ -1,11 +1,11 @@
-from khaosz.inference.core import (
+from astrai.inference.core import (
     disable_random_init,
     GeneratorCore,
     EmbeddingEncoderCore,
     KVCacheManager,
 )
 
-from khaosz.inference.generator import (
+from astrai.inference.generator import (
     GenerationRequest,
     LoopGenerator,
     StreamGenerator,
diff --git a/khaosz/inference/core.py b/astrai/inference/core.py
similarity index 99%
rename from khaosz/inference/core.py
rename to astrai/inference/core.py
index 77129ce..db3a433 100644
--- a/khaosz/inference/core.py
+++ b/astrai/inference/core.py
@@ -4,7 +4,7 @@ import torch.nn as nn
 from torch import Tensor
 from contextlib import contextmanager
 from typing import Any, Callable, List, Tuple, Union, Optional, Self
-from khaosz.config import ModelParameter, ModelConfig
+from astrai.config import ModelParameter, ModelConfig
 
 
 def apply_sampling_strategies(
diff --git a/khaosz/inference/generator.py b/astrai/inference/generator.py
similarity index 98%
rename from khaosz/inference/generator.py
rename to astrai/inference/generator.py
index 534ac23..89d10f3 100644
--- a/khaosz/inference/generator.py
+++ b/astrai/inference/generator.py
@@ -2,8 +2,8 @@ import torch
 from dataclasses import dataclass
 from torch import Tensor
 from typing import List, Tuple, Union, Optional, Generator
-from khaosz.inference.core import GeneratorCore, EmbeddingEncoderCore, KVCacheManager
-from khaosz.config.param_config import ModelParameter
+from astrai.inference.core import GeneratorCore, EmbeddingEncoderCore, KVCacheManager
+from astrai.config.param_config import ModelParameter
 
 
 HistoryType = List[Tuple[str, str]]
diff --git a/khaosz/model/__init__.py b/astrai/model/__init__.py
similarity index 62%
rename from khaosz/model/__init__.py
rename to astrai/model/__init__.py
index 4f71990..dccfdc3 100644
--- a/khaosz/model/__init__.py
+++ b/astrai/model/__init__.py
@@ -1,10 +1,10 @@
-from khaosz.model.module import (
+from astrai.model.module import (
     Linear,
     RMSNorm,
     MLP,
     GQA,
     DecoderBlock,
 )
-from khaosz.model.transformer import Transformer
+from astrai.model.transformer import Transformer
 
 __all__ = ["Linear", "RMSNorm", "MLP", "GQA", "DecoderBlock", "Transformer"]
diff --git a/khaosz/model/module.py b/astrai/model/module.py
similarity index 100%
rename from khaosz/model/module.py
rename to astrai/model/module.py
diff --git a/khaosz/model/transformer.py b/astrai/model/transformer.py
similarity index 98%
rename from khaosz/model/transformer.py
rename to astrai/model/transformer.py
index a90fade..0608d23 100644
--- a/khaosz/model/transformer.py
+++ b/astrai/model/transformer.py
@@ -3,8 +3,8 @@ import torch.nn as nn
 
 from torch import Tensor
 from typing import Any, Mapping, Optional, Tuple
-from khaosz.config.model_config import ModelConfig
-from khaosz.model.module import (
+from astrai.config.model_config import ModelConfig
+from astrai.model.module import (
     Embedding,
     DecoderBlock,
     Linear,
diff --git a/khaosz/parallel/__init__.py b/astrai/parallel/__init__.py
similarity index 76%
rename from khaosz/parallel/__init__.py
rename to astrai/parallel/__init__.py
index 896be8a..562eff4 100644
--- a/khaosz/parallel/__init__.py
+++ b/astrai/parallel/__init__.py
@@ -1,4 +1,4 @@
-from khaosz.parallel.setup import (
+from astrai.parallel.setup import (
     get_world_size,
     get_rank,
     get_current_device,
@@ -7,7 +7,7 @@ from khaosz.parallel.setup import (
     spawn_parallel_fn,
 )
 
-from khaosz.parallel.module import RowParallelLinear, ColumnParallelLinear
+from astrai.parallel.module import RowParallelLinear, ColumnParallelLinear
 
 __all__ = [
     "get_world_size",
diff --git a/khaosz/parallel/module.py b/astrai/parallel/module.py
similarity index 100%
rename from khaosz/parallel/module.py
rename to astrai/parallel/module.py
diff --git a/khaosz/parallel/setup.py b/astrai/parallel/setup.py
similarity index 100%
rename from khaosz/parallel/setup.py
rename to astrai/parallel/setup.py
diff --git a/khaosz/trainer/__init__.py b/astrai/trainer/__init__.py
similarity index 70%
rename from khaosz/trainer/__init__.py
rename to astrai/trainer/__init__.py
index b266a66..6165d57 100644
--- a/khaosz/trainer/__init__.py
+++ b/astrai/trainer/__init__.py
@@ -1,8 +1,8 @@
-from khaosz.trainer.trainer import Trainer
-from khaosz.trainer.strategy import StrategyFactory, BaseStrategy
-from khaosz.trainer.schedule import SchedulerFactory, BaseScheduler
+from astrai.trainer.trainer import Trainer
+from astrai.trainer.strategy import StrategyFactory, BaseStrategy
+from astrai.trainer.schedule import SchedulerFactory, BaseScheduler
 
-from khaosz.trainer.train_callback import (
+from astrai.trainer.train_callback import (
     TrainCallback,
     GradientClippingCallback,
     SchedulerCallback,
diff --git a/khaosz/trainer/metric_util.py b/astrai/trainer/metric_util.py
similarity index 100%
rename from khaosz/trainer/metric_util.py
rename to astrai/trainer/metric_util.py
diff --git a/khaosz/trainer/schedule.py b/astrai/trainer/schedule.py
similarity index 99%
rename from khaosz/trainer/schedule.py
rename to astrai/trainer/schedule.py
index 453d2ee..f0eacbd 100644
--- a/khaosz/trainer/schedule.py
+++ b/astrai/trainer/schedule.py
@@ -4,7 +4,7 @@ import math
 from abc import abstractmethod, ABC
 from typing import Any, Dict, List, Type
 from torch.optim.lr_scheduler import LRScheduler
-from khaosz.config.schedule_config import ScheduleConfig
+from astrai.config.schedule_config import ScheduleConfig
 
 
 class BaseScheduler(LRScheduler, ABC):
diff --git a/khaosz/trainer/strategy.py b/astrai/trainer/strategy.py
similarity index 100%
rename from khaosz/trainer/strategy.py
rename to astrai/trainer/strategy.py
diff --git a/khaosz/trainer/train_callback.py b/astrai/trainer/train_callback.py
similarity index 97%
rename from khaosz/trainer/train_callback.py
rename to astrai/trainer/train_callback.py
index 08de779..997cd48 100644
--- a/khaosz/trainer/train_callback.py
+++ b/astrai/trainer/train_callback.py
@@ -8,8 +8,8 @@ from tqdm import tqdm
 from torch.nn.utils import clip_grad_norm_
 from typing import Callable, List, Optional, Protocol
 
-from khaosz.parallel import only_on_rank
-from khaosz.trainer.metric_util import (
+from astrai.parallel import only_on_rank
+from astrai.trainer.metric_util import (
     ctx_get_loss,
     ctx_get_lr,
     ctx_get_grad_max,
@@ -19,8 +19,8 @@ from khaosz.trainer.metric_util import (
     ctx_get_grad_std,
     ctx_get_grad_nan_num,
 )
-from khaosz.data.serialization import Checkpoint
-from khaosz.trainer.train_context import TrainContext
+from astrai.data.serialization import Checkpoint
+from astrai.trainer.train_context import TrainContext
 
 
 class TrainCallback(Protocol):
diff --git a/khaosz/trainer/train_context.py b/astrai/trainer/train_context.py
similarity index 91%
rename from khaosz/trainer/train_context.py
rename to astrai/trainer/train_context.py
index f718ffd..7f65b02 100644
--- a/khaosz/trainer/train_context.py
+++ b/astrai/trainer/train_context.py
@@ -3,11 +3,11 @@ from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 
-from khaosz.data import ResumableDistributedSampler
-from khaosz.data.serialization import Checkpoint
-from khaosz.trainer.strategy import StrategyFactory, BaseStrategy
-from khaosz.config.train_config import TrainConfig
-from khaosz.parallel.setup import get_current_device, get_world_size, get_rank
+from astrai.data import ResumableDistributedSampler
+from astrai.data.serialization import Checkpoint
+from astrai.trainer.strategy import StrategyFactory, BaseStrategy
+from astrai.config.train_config import TrainConfig
+from astrai.parallel.setup import get_current_device, get_world_size, get_rank
 
 from dataclasses import dataclass, field
 from typing import Optional, Self
diff --git a/khaosz/trainer/trainer.py b/astrai/trainer/trainer.py
similarity index 93%
rename from khaosz/trainer/trainer.py
rename to astrai/trainer/trainer.py
index 5a5a88e..74c35a6 100644
--- a/khaosz/trainer/trainer.py
+++ b/astrai/trainer/trainer.py
@@ -1,7 +1,7 @@
 import logging
 from typing import Optional, List
-from khaosz.config import TrainConfig
-from khaosz.trainer.train_callback import (
+from astrai.config import TrainConfig
+from astrai.trainer.train_callback import (
     TrainCallback,
     ProgressBarCallback,
     CheckpointCallback,
@@ -9,9 +9,9 @@ from khaosz.trainer.train_callback import (
     GradientClippingCallback,
     SchedulerCallback,
 )
-from khaosz.trainer.train_context import TrainContext, TrainContextBuilder
-from khaosz.data.serialization import Checkpoint
-from khaosz.parallel.setup import spawn_parallel_fn
+from astrai.trainer.train_context import TrainContext, TrainContextBuilder
+from astrai.data.serialization import Checkpoint
+from astrai.parallel.setup import spawn_parallel_fn
 
 logger = logging.getLogger(__name__)
 
diff --git a/demo/download.py b/demo/download.py
index 8cb9052..fb3150f 100644
--- a/demo/download.py
+++ b/demo/download.py
@@ -6,7 +6,7 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 if __name__ == "__main__":
     snapshot_download(
-        repo_id="ViperEk/KHAOSZ",
+        repo_id="ViperEk/AstrAI",
         local_dir=PARAMETER_ROOT,
         force_download=True,
     )
diff --git a/demo/generate_ar.py b/demo/generate_ar.py
index 87033d0..c75a601 100644
--- a/demo/generate_ar.py
+++ b/demo/generate_ar.py
@@ -1,8 +1,8 @@
 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 GeneratorFactory, GenerationRequest
+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
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
diff --git a/demo/generate_batch.py b/demo/generate_batch.py
index 0bf7645..4edd88e 100644
--- a/demo/generate_batch.py
+++ b/demo/generate_batch.py
@@ -1,8 +1,8 @@
 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 GeneratorFactory, GenerationRequest
+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
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
diff --git a/demo/stream_chat.py b/demo/stream_chat.py
index d119685..fbe31ea 100644
--- a/demo/stream_chat.py
+++ b/demo/stream_chat.py
@@ -1,8 +1,8 @@
 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 GeneratorFactory, GenerationRequest
+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
 PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
diff --git a/pyproject.toml b/pyproject.toml
index bd0a69c..850cd40 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 dynamic = ["version"]
-name = "khaosz"
+name = "astrai"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
@@ -23,7 +23,7 @@ classifiers = [
     "License :: OSI Approved :: GPL-3.0",
     "Operating System :: OS Independent",
 ]
-urls = { Homepage = "https://github.com/ViperEkura/KHAOSZ" }
+urls = { Homepage = "https://github.com/ViperEkura/AstrAI" }
 
 [project.optional-dependencies]
 dev = ["pytest==9.0.2", "ruff"]
@@ -35,7 +35,7 @@ where = ["."]
 extra-index-url = "https://download.pytorch.org/whl/cu126"
 
 [tool.setuptools.dynamic]
-version = { attr = "khaosz.__version__" }
+version = { attr = "astrai.__version__" }
 
 [tool.ruff]
 target-version = "py312"
diff --git a/tests/conftest.py b/tests/conftest.py
index b21fa92..ffe8a8a 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -7,9 +7,9 @@ import torch
 import pytest
 
 from torch.utils.data import Dataset
-from khaosz.config.model_config import ModelConfig
-from khaosz.data.tokenizer import BpeTokenizer
-from khaosz.model.transformer import Transformer
+from astrai.config.model_config import ModelConfig
+from astrai.data.tokenizer import BpeTokenizer
+from astrai.model.transformer import Transformer
 
 
 class RandomDataset(Dataset):
diff --git a/tests/data/test_checkpoint.py b/tests/data/test_checkpoint.py
index abb9b33..33ccb66 100644
--- a/tests/data/test_checkpoint.py
+++ b/tests/data/test_checkpoint.py
@@ -4,8 +4,8 @@ import torch.distributed as dist
 
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from khaosz.data.serialization import Checkpoint
-from khaosz.parallel.setup import get_rank, spawn_parallel_fn
+from astrai.data.serialization import Checkpoint
+from astrai.parallel.setup import get_rank, spawn_parallel_fn
 
 
 def test_single_process():
diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py
index 0420e23..27a077b 100644
--- a/tests/data/test_dataset.py
+++ b/tests/data/test_dataset.py
@@ -1,8 +1,8 @@
 import torch
 import numpy as np
 
-from khaosz.data.serialization import save_h5
-from khaosz.data.dataset import *
+from astrai.data.serialization import save_h5
+from astrai.data.dataset import *
 
 
 def test_dataset_loader_random_paths(base_test_env):
diff --git a/tests/data/test_sampler.py b/tests/data/test_sampler.py
index 94f7a77..2c461c1 100644
--- a/tests/data/test_sampler.py
+++ b/tests/data/test_sampler.py
@@ -1,5 +1,5 @@
-from khaosz.trainer import *
-from khaosz.data import *
+from astrai.trainer import *
+from astrai.data import *
 
 
 def test_random_sampler_consistency(random_dataset):
diff --git a/tests/module/test_module.py b/tests/module/test_module.py
index 03c60b3..a37bbc1 100644
--- a/tests/module/test_module.py
+++ b/tests/module/test_module.py
@@ -5,11 +5,11 @@ import shutil
 import pytest
 import tempfile
 import safetensors.torch as st
-from khaosz.trainer import *
-from khaosz.config import *
-from khaosz.model import *
-from khaosz.data import *
-from khaosz.inference.generator import EmbeddingEncoderCore, GeneratorCore
+from astrai.trainer import *
+from astrai.config import *
+from astrai.model import *
+from astrai.data import *
+from astrai.inference.generator import EmbeddingEncoderCore, GeneratorCore
 from tokenizers import pre_tokenizers
 
 
diff --git a/tests/module/test_tie_weight.py b/tests/module/test_tie_weight.py
index 221f180..b628aca 100644
--- a/tests/module/test_tie_weight.py
+++ b/tests/module/test_tie_weight.py
@@ -4,8 +4,8 @@ import torch
 import pytest
 import tempfile
 import safetensors.torch as st
-from khaosz.model.transformer import Transformer
-from khaosz.config.model_config import ModelConfig
+from astrai.model.transformer import Transformer
+from astrai.config.model_config import ModelConfig
 
 
 @pytest.fixture
diff --git a/tests/test_parallel.py b/tests/test_parallel.py
index c5207b9..c05b084 100644
--- a/tests/test_parallel.py
+++ b/tests/test_parallel.py
@@ -1,7 +1,7 @@
 import torch
 import torch.distributed as dist
 
-from khaosz.parallel import get_rank, only_on_rank, spawn_parallel_fn
+from astrai.parallel import get_rank, only_on_rank, spawn_parallel_fn
 
 
 @only_on_rank(0)
diff --git a/tests/trainer/test_callbacks.py b/tests/trainer/test_callbacks.py
index 08fef7b..58eda2e 100644
--- a/tests/trainer/test_callbacks.py
+++ b/tests/trainer/test_callbacks.py
@@ -1,7 +1,7 @@
 import torch
 
-from khaosz.config import *
-from khaosz.trainer import *
+from astrai.config import *
+from astrai.trainer import *
 
 
 def test_callback_integration(base_test_env, random_dataset):
diff --git a/tests/trainer/test_early_stopping.py b/tests/trainer/test_early_stopping.py
index a9a244a..a530cdb 100644
--- a/tests/trainer/test_early_stopping.py
+++ b/tests/trainer/test_early_stopping.py
@@ -1,9 +1,9 @@
 import os
 import torch
 import numpy as np
-from khaosz.config import *
-from khaosz.trainer import *
-from khaosz.data.serialization import Checkpoint
+from astrai.config import *
+from astrai.trainer import *
+from astrai.data.serialization import Checkpoint
 
 
 def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
diff --git a/tests/trainer/test_train_strategy.py b/tests/trainer/test_train_strategy.py
index ac7ad30..339476e 100644
--- a/tests/trainer/test_train_strategy.py
+++ b/tests/trainer/test_train_strategy.py
@@ -2,9 +2,9 @@ import torch
 import numpy as np
 import pytest
 
-from khaosz.config import *
-from khaosz.trainer.schedule import *
-from khaosz.data.dataset import *
+from astrai.config import *
+from astrai.trainer.schedule import *
+from astrai.data.dataset import *
 
 
 def test_schedule_factory_random_configs():
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index a0cca66..3f69398 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -2,9 +2,9 @@ import torch
 import numpy as np
 
 
-from khaosz.config import *
-from khaosz.trainer import *
-from khaosz.data.dataset import *
+from astrai.config import *
+from astrai.trainer import *
+from astrai.data.dataset import *
 
 
 def test_different_batch_sizes(base_test_env, random_dataset):
diff --git a/tools/benchmark.py b/tools/benchmark.py
index ebfd7af..211cef5 100644
--- a/tools/benchmark.py
+++ b/tools/benchmark.py
@@ -1,7 +1,7 @@
 import torch
 from typing import Dict, Any
 from dataclasses import dataclass
-from khaosz.model.transformer import ModelConfig, Transformer
+from astrai.model.transformer import ModelConfig, Transformer
 
 
 @dataclass
diff --git a/tools/generate.py b/tools/generate.py
index 22f1dc6..2f78b2a 100644
--- a/tools/generate.py
+++ b/tools/generate.py
@@ -2,9 +2,9 @@ import torch
 import json
 import argparse
 
-from khaosz.config.param_config import ModelParameter
-from khaosz.inference.generator import BatchGenerator, GenerationRequest
-from khaosz.inference.core import disable_random_init
+from astrai.config.param_config import ModelParameter
+from astrai.inference.generator import BatchGenerator, GenerationRequest
+from astrai.inference.core import disable_random_init
 
 
 def processor(
diff --git a/tools/perplexity.py b/tools/perplexity.py
index a91d370..e6a8de8 100644
--- a/tools/perplexity.py
+++ b/tools/perplexity.py
@@ -6,8 +6,8 @@ import argparse
 import tqdm
 
 from torch import Tensor
-from khaosz.config.param_config import ModelParameter
-from khaosz.inference.core import disable_random_init
+from astrai.config.param_config import ModelParameter
+from astrai.inference.core import disable_random_init
 
 
 def compute_perplexity(
diff --git a/tools/train.py b/tools/train.py
index 69795f7..86260ed 100644
--- a/tools/train.py
+++ b/tools/train.py
@@ -6,10 +6,10 @@ import torch.optim as optim
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from functools import partial
-from khaosz.data import DatasetLoader
-from khaosz.config import ModelParameter, TrainConfig, CosineScheduleConfig
-from khaosz.trainer import Trainer, SchedulerFactory
-from khaosz.parallel import get_rank
+from astrai.data import DatasetLoader
+from astrai.config import ModelParameter, TrainConfig, CosineScheduleConfig
+from astrai.trainer import Trainer, SchedulerFactory
+from astrai.parallel import get_rank
 
 
 def parse_args() -> argparse.Namespace: