# SPDX-License-Identifier: Apache-2.0

from argparse import Namespace
import os
import numpy as np

from vllm import LLM, EngineArgs
from vllm.utils import FlexibleArgumentParser


def main(args: Namespace):

    prompts = [
        "你好，我的名字是",
        "美国总统是",
        "法国的首都是",
        "人工智能的未来是",
    ]


    model = LLM(**vars(args))

    outputs = model.embed(prompts)
    script_dir = os.path.dirname(os.path.abspath(__file__))

    embeddings = [output.outputs.embedding for output in outputs]

    output_path = os.path.join(script_dir, 'embeddings_A800.npy')
    np.save(output_path, np.array(embeddings))
    

    for i, (prompt, output) in enumerate(zip(prompts, outputs)):
        embeds = output.outputs.embedding
        embeds_trimmed = ((str(embeds[:16])[:-1] +
                          ", ...]") if len(embeds) > 16 else embeds)
        print(f"提示: {prompt!r} | "
              f"嵌入: {embeds_trimmed} (大小={len(embeds)})")
    print(f"所有嵌入已保存到: {output_path}")


if __name__ == "__main__":
    parser = FlexibleArgumentParser()
    parser = EngineArgs.add_cli_args(parser)
  
    parser.set_defaults(model="intfloat/multilingual-e5-large",
                       task="embed",
                       enforce_eager=True)
    args = parser.parse_args()
    main(args)