Files
vllm-omni/examples/offline_inference/tmp_test/extract_prompts.py
T
wzliu 33e03b44c5 added omni mooncake connectors to distributed directory, integrate omni connector to omni-llm, rebased with main (twice).
ray + mooncake connector work on 2 nodes, but needes several updates\n1.start ray inside vllm-omni \n2.placement group optimize \n3.fix close
2025-12-01 10:45:26 +08:00

46 lines
1.2 KiB
Python

#!/usr/bin/env python3
import argparse
from typing import Optional
def extract_prompt(line: str) -> Optional[str]:
# 提取第一个 '|' 与第二个 '|' 之间的内容
i = line.find("|")
if i == -1:
return None
j = line.find("|", i + 1)
if j == -1:
return None
return line[i + 1 : j].strip()
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--input", "-i", required=True, help="Input .lst file path")
parser.add_argument("--output", "-o", required=True, help="Output file path")
parser.add_argument(
"--topk",
"-k",
type=int,
default=100,
help="Extract the top K prompts (default: 100)",
)
args = parser.parse_args()
prompts = []
with open(args.input, encoding="utf-8", errors="ignore") as f:
for line in f:
if len(prompts) >= args.topk:
break
p = extract_prompt(line.rstrip("\n"))
if p:
prompts.append(p)
with open(args.output, "w", encoding="utf-8") as f:
for p in prompts:
f.write(p + "\n")
if __name__ == "__main__":
main()