Function main

examples/llm-api/_tensorrt_engine/llm_quantization.py:49–76

()

def main():

    for quant_config, calib_config in quant_and_calib_configs:
        # The built-in end-to-end quantization is triggered according to the passed quant_config.
        llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
                  quant_config=quant_config,
                  calib_config=calib_config)

        # Sample prompts.
        prompts = [
            "Hello, my name is",
            "The capital of France is",
            "The future of AI is",
        ]

        # Create a sampling params.
        sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

        for output in llm.generate(prompts, sampling_params):
            print(
                f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
            )
        llm.shutdown()

    # Got output like
    # Prompt: 'Hello, my name is', Generated text: 'Jane Smith. I am a resident of the city. Can you tell me more about the public services provided in the area?'
    # Prompt: 'The capital of France is', Generated text: 'located in Paris, France. The population of Paris, France, is estimated to be 2 million. France is home to many famous artists, including Picasso'
    # Prompt: 'The future of AI is', Generated text: 'an open and collaborative project. The project is an ongoing effort, and we invite participation from members of the community.\n\nOur community is'


if __name__ == '__main__':