main
py 50 lines 1.58 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from vllm import LLM, SamplingParams
5 from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Vector
6
7 # Sample prompts.
8 prompts = [
9 "Hello, my name is",
10 "The president of the United States is",
11 "The capital of France is",
12 "The future of AI is",
13 ]
14 # Create a sampling params object.
15 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
16
17
18 def main():
19 # Create an LLM.
20 llm = LLM(model="facebook/opt-125m", disable_log_stats=False)
21
22 # Generate texts from the prompts.
23 outputs = llm.generate(prompts, sampling_params)
24
25 # Print the outputs.
26 print("-" * 50)
27 for output in outputs:
28 prompt = output.prompt
29 generated_text = output.outputs[0].text
30 print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
31 print("-" * 50)
32
33 # Dump all metrics
34 for metric in llm.get_metrics():
35 if isinstance(metric, Gauge):
36 print(f"{metric.name} (gauge) = {metric.value}")
37 elif isinstance(metric, Counter):
38 print(f"{metric.name} (counter) = {metric.value}")
39 elif isinstance(metric, Vector):
40 print(f"{metric.name} (vector) = {metric.values}")
41 elif isinstance(metric, Histogram):
42 print(f"{metric.name} (histogram)")
43 print(f" sum = {metric.sum}")
44 print(f" count = {metric.count}")
45 for bucket_le, value in metric.buckets.items():
46 print(f" {bucket_le} = {value}")
47
48
49 if __name__ == "__main__":
50 main()