| import sglang as sgl | |
| def main(): | |
| # 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. | |
| llm = sgl.Engine( | |
| model_path="Alibaba-NLP/gte-Qwen2-1.5B-instruct", is_embedding=True | |
| ) | |
| outputs = llm.encode(prompts) | |
| # Print the outputs (embedding vectors) | |
| for prompt, output in zip(prompts, outputs): | |
| print("===============================") | |
| print(f"Prompt: {prompt}\nEmbedding vector: {output['embedding']}") | |
| # The __main__ condition is necessary here because we use "spawn" to create subprocesses | |
| # Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine | |
| if __name__ == "__main__": | |
| main() | |