refactor: 优化工具脚本接口并修复批处理问题

This commit is contained in:
2026-04-05 21:56:22 +08:00
parent e58dbd7c57
commit 9b22b1651e
8 changed files with 140 additions and 64 deletions
+1 -3
View File
@@ -13,11 +13,9 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
def generate_text():
# Load model from pretrained
model = AutoModel.from_pretrained(PARAMETER_ROOT)
tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
model.to(device="cuda", dtype=torch.bfloat16)
# Load tokenizer from pretrained
tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT / "tokenizer")
query = input(">> ")
engine = InferenceEngine(