chore: 简化格式并更新文档

khaosz/data/dataset.py

@@ -7,7 +7,7 @@ from abc import ABC, abstractmethod
 from torch import Tensor
 from torch.utils.data import Dataset
 from khaosz.data.serialization import load_h5
-from typing import Callable, List, Dict, Literal, Optional, Union
+from typing import List, Dict, Optional, Union
 
 
 class BaseSegmentFetcher:

khaosz/data/serialization.py

@@ -75,7 +75,7 @@ class Checkpoint:
             with open(save_path / "meta.json", "w") as f:
                 json.dump(meta, f, indent=2)
 
-            st.save_file(self.state_dict, save_path / f"state_dict.safetensors")
+            st.save_file(self.state_dict, save_path / "state_dict.safetensors")
 
     @classmethod
     def load(
@@ -96,7 +96,7 @@ class Checkpoint:
             dist.broadcast_object_list(meta_list, src=0)
             meta = meta_list[0]
 
-        state_dict = st.load_file(save_path / f"state_dict.safetensors")
+        state_dict = st.load_file(save_path / "state_dict.safetensors")
 
         return cls(
             state_dict=state_dict,

khaosz/inference/generator.py

@@ -219,7 +219,7 @@ class BatchGenerator(GeneratorCore):
                     ids_list[i].append(token)
                     c_ids += 1
 
-                    is_active = not token in self.tokenizer.stop_ids
+                    is_active = token not in self.tokenizer.stop_ids
                     activate_task_mask[i] = is_active
                     active_mask.append(is_active)
 

khaosz/trainer/strategy.py

@@ -7,7 +7,7 @@ import torch.nn.functional as F
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from torch import Tensor
-from typing import Any, Callable, Dict, Union, Optional
+from typing import Any, Callable, Dict, Union
 from abc import ABC, abstractmethod
 
 

khaosz/trainer/train_callback.py

@@ -6,7 +6,6 @@ import torch.nn as nn
 from pathlib import Path
 from tqdm import tqdm
 from torch.nn.utils import clip_grad_norm_
-from torch.optim.lr_scheduler import LRScheduler
 from typing import Callable, List, Optional, Protocol
 
 from khaosz.parallel import only_on_rank