refactor: decode 按页分桶批处理,position_ids 改为 per-task 构建

This commit is contained in:
2026-05-14 14:22:11 +08:00
parent c0effc9f5b
commit 6269bacfc3
3 changed files with 16 additions and 13 deletions
+2 -1
View File
@@ -18,6 +18,7 @@ def processor(
question_key: str,
response_key: str,
max_tokens: int,
batch_size: int,
):
# Load model and tokenizer
model = AutoModel.from_pretrained(param_path)
@@ -25,7 +26,7 @@ def processor(
model.to(device="cuda", dtype=torch.bfloat16)
# Create inference engine
engine = InferenceEngine(model=model, tokenizer=tokenizer)
engine = InferenceEngine(model=model, tokenizer=tokenizer, max_batch_size=batch_size)
with open(input_json_file, "r", encoding="utf-8") as f:
input_data = [json.loads(line) for line in f]