| """批量翻译 —— vLLM 后端。 |
| |
| pip install vllm |
| python example_vllm.py |
| |
| **`vllm serve` 开箱用不了。** piece 不是 HF 分词器,vLLM 起不来它的 |
| tokenizer,所以必须: |
| |
| 1. `LLM(..., skip_tokenizer_init=True)` —— 不让 vLLM 自己加载分词器 |
| 2. 提示词自己编码,用 `TokensPrompt(prompt_token_ids=...)` 喂进去 |
| 3. 输出拿 `token_ids` 再用本目录的 `tokenizer.py` 解码 |
| |
| 这也是训练时评测走的路径,数字可比。 |
| """ |
| import os |
| import sys |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from tokenizer import PieceTokenizerWrapper |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
|
|
| PROMPT = { |
| "zh-en": "Translate the following text from Chinese to English.\n" |
| "Chinese: {src}\nEnglish:", |
| "en-zh": "Translate the following text from English to Chinese.\n" |
| "English: {src}\nChinese:", |
| } |
|
|
|
|
| def main(): |
| from vllm import LLM, SamplingParams |
| from vllm.inputs import TokensPrompt |
|
|
| tok = PieceTokenizerWrapper(HERE) |
| llm = LLM(model=HERE, skip_tokenizer_init=True, dtype="bfloat16", |
| max_model_len=1024, trust_remote_code=True) |
|
|
| batch = [("人工智能正在深刻改变我们的生活方式。", "zh-en"), |
| ("机器翻译的基本任务是把一种语言转换成另一种语言。", "zh-en"), |
| ("The quick brown fox jumps over the lazy dog.", "en-zh"), |
| ("Machine translation has improved dramatically in recent years.", |
| "en-zh")] |
|
|
| prompts = [TokensPrompt(prompt_token_ids=tok.apply_chat_template( |
| [{"role": "user", "content": PROMPT[d].format(src=s)}], |
| tokenize=True, add_generation_prompt=True)) for s, d in batch] |
|
|
| |
| params = SamplingParams(max_tokens=256, temperature=0.0, |
| stop_token_ids=[tok.eos_token_id]) |
|
|
| for (src, d), out in zip(batch, llm.generate(prompts, params)): |
| ids = list(out.outputs[0].token_ids) |
| if ids and ids[-1] == tok.eos_token_id: |
| ids = ids[:-1] |
| print(f"[{d}] {src}\n → {tok.decode(ids, skip_special_tokens=True).strip()}\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
