chore: 修改类名，优化导入顺序

astrai/inference/engine.py

@@ -8,8 +8,8 @@ from typing import Any, Dict, Generator, List, Optional, Union
 import torch
 import torch.nn as nn
 
-from astrai.tokenize.tokenizer import TextTokenizer
 from astrai.inference.scheduler import InferenceScheduler
+from astrai.tokenize import AutoTokenizer
 
 logger = logging.getLogger(__name__)
 
@@ -109,7 +109,7 @@ class InferenceEngine:
     def __init__(
         self,
         model: nn.Module,
-        tokenizer: TextTokenizer,
+        tokenizer: AutoTokenizer,
         max_batch_size: int = 1,
         max_seq_len: Optional[int] = None,
     ):

astrai/inference/scheduler.py

@@ -9,7 +9,7 @@ import torch
 from torch import Tensor
 
 from astrai.model.automodel import AutoModel
-from astrai.tokenize.tokenizer import TextTokenizer
+from astrai.tokenize import AutoTokenizer
 
 
 class TaskStatus:
@@ -101,7 +101,7 @@ class InferenceScheduler:
     def __init__(
         self,
         model: AutoModel,
-        tokenizer: TextTokenizer,
+        tokenizer: AutoTokenizer,
         max_batch_size: int = 16,
         max_seq_len: Optional[int] = None,
         device: str = "cuda",

astrai/inference/server.py

@@ -19,7 +19,7 @@ from pydantic import BaseModel, Field
 
 from astrai.inference.engine import InferenceEngine
 from astrai.model import AutoModel
-from astrai.tokenize import TextTokenizer
+from astrai.tokenize import AutoTokenizer
 
 logger = logging.getLogger(__name__)
 
@@ -96,7 +96,7 @@ def load_model(
         raise FileNotFoundError(f"Parameter directory not found: {param_path}")
 
     # Load tokenizer separately
-    tokenizer = TextTokenizer.from_pretrained(param_path)
+    tokenizer = AutoTokenizer.from_pretrained(param_path)
     _model_param = AutoModel.from_pretrained(param_path)
     _model_param.to(device=device, dtype=dtype)
     logger.info(f"Model loaded on {device} with dtype {dtype}")