2024-05-11 11:30:37 -07:00
|
|
|
from vllm import LLM
|
|
|
|
|
|
|
|
# Sample prompts.
|
|
|
|
prompts = [
|
|
|
|
"Hello, my name is",
|
|
|
|
"The president of the United States is",
|
|
|
|
"The capital of France is",
|
|
|
|
"The future of AI is",
|
|
|
|
]
|
|
|
|
|
|
|
|
# Create an LLM.
|
2024-12-11 17:28:00 +08:00
|
|
|
model = LLM(
|
|
|
|
model="intfloat/e5-mistral-7b-instruct",
|
|
|
|
task="embed", # You should pass task="embed" for embedding models
|
|
|
|
enforce_eager=True,
|
|
|
|
)
|
|
|
|
|
2024-12-01 14:36:51 +08:00
|
|
|
# Generate embedding. The output is a list of PoolingRequestOutputs.
|
2024-05-11 11:30:37 -07:00
|
|
|
outputs = model.encode(prompts)
|
|
|
|
# Print the outputs.
|
|
|
|
for output in outputs:
|
|
|
|
print(output.outputs.embedding) # list of 4096 floats
|