#!/usr/bin/env python
# coding: utf-8

# <a href="https://colab.research.google.com/github/jerryjliu/llama_index/blob/main/docs/docs/examples/llm/nvidia_tensorrt.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="在 Colab 中打开"/></a>
# 

# Nvidia TensorRT-LLM是一种用于低延迟推理的深度学习模型优化工具。TensorRT-LLM使用Nvidia的TensorRT推理引擎，通过减少模型推理的延迟来提高性能。它还提供了一些优化技术，如层次化的量化和剪枝，以进一步优化模型。TensorRT-LLM旨在帮助开发人员在边缘设备上部署实时的深度学习推理应用程序。
# 

# TensorRT-LLM为用户提供了一个易于使用的Python API，用于定义大型语言模型（LLMs）并构建包含最先进优化的TensorRT引擎，以在NVIDIA GPU上高效进行推理。
# 
# [TensorRT-LLM Github](https://github.com/NVIDIA/TensorRT-LLM)
# 

# ## TensorRT-LLM环境设置
# 由于TensorRT-LLM是一个用于与本地模型进行进程交互的SDK，因此必须遵循一些环境设置步骤，以确保可以使用TensorRT-LLM设置。请注意，目前需要Nvidia Cuda 12.2或更高版本才能运行TensorRT-LLM。
# 
# 在本教程中，我们将展示如何将连接器与GPT2模型一起使用。
# 为了获得最佳体验，我们建议按照官方[TensorRT-LLM Github](https://github.com/NVIDIA/TensorRT-LLM)上的[安装](https://github.com/NVIDIA/TensorRT-LLM/tree/v0.8.0?tab=readme-ov-file#installation)过程进行操作。
# 
# 以下步骤展示了如何为x86_64用户设置您的模型与TensorRT-LLM v0.8.0。
# 
# 1. 获取并启动基本的docker镜像环境。
# ```
# docker run --rm --runtime=nvidia --gpus all --entrypoint /bin/bash -it nvidia/cuda:12.1.0-devel-ubuntu22.04
# ```
# 
# 2. 安装依赖项，TensorRT-LLM需要Python 3.10
# ```
# apt-get update && apt-get -y install python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs wget
# ```
# 3. 安装最新稳定版本（对应发布分支）的TensorRT-LLM。我们使用的是版本0.8.0，但要获取最新版本，请参考[官方发布页面](https://github.com/NVIDIA/TensorRT-LLM/releases)。
# ```
# pip3 install tensorrt_llm==0.8.0 -U --extra-index-url https://pypi.nvidia.com
# ```
# 
# 4. 检查安装
# ```
# python3 -c "import tensorrt_llm"
# ```
# 上述命令不应产生任何错误。
# 
# 5. 对于本示例，我们将使用GPT2。GPT2模型文件需要通过脚本按照[这里](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/gpt#usage)的说明创建。
#     * 首先，在第1阶段启动的容器内，克隆TensorRT-LLM存储库：
#     ```
#     git clone --branch v0.8.0 https://github.com/NVIDIA/TensorRT-LLM.git
#     ```
#     * 使用以下命令安装GPT2模型的要求：
#     ```
#     cd TensorRT-LLM/examples/gpt/ && pip install -r requirements.txt
#     ```
#     * 下载hf gpt2模型
#     ```
#     rm -rf gpt2 && git clone https://huggingface.co/gpt2-medium gpt2
#     cd gpt2
#     rm pytorch_model.bin model.safetensors
#     wget -q https://huggingface.co/gpt2-medium/resolve/main/pytorch_model.bin
#     cd ..
#     ```
#     * 将权重从HF Transformers转换为TensorRT-LLM格式
#     ```
#     python3 hf_gpt_convert.py -i gpt2 -o ./c-model/gpt2 --tensor-parallelism 1 --storage-type float16
#     ```
#     * 构建TensorRT引擎
#     ```
#     python3 build.py --model_dir=./c-model/gpt2/1-gpu --use_gpt_attention_plugin --remove_input_padding
#     ```
#   
# 6. 安装`llama-index-llms-nvidia-tensorrt`包
#   ```
#   pip install llama-index-llms-nvidia-tensorrt
#   ```
# 

# ## 基本用法
# 

# #### 使用提示调用`complete`
# 

# ```python
# from llama_index.llms.nvidia_tensorrt import LocalTensorRTLLM
# 
# llm = LocalTensorRTLLM(
#     model_path="./engine_outputs",
#     engine_name="gpt_float16_tp1_rank0.engine",
#     tokenizer_dir="gpt2",
#     max_new_tokens=40,
# )
# 
# resp = llm.complete("谁是哈利波特？")
# print(str(resp))
# ```
# 

# 哈利·波特是J.K.罗琳在她的第一部小说《哈利·波特与魔法石》中创造的虚构人物。这个角色是一个生活在虚构小镇的巫师。
#