code: # SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  # SPDX-License-Identifier: MIT
  #
  # Permission is hereby granted, free of charge, to any person obtaining a
  # copy of this software and associated documentation files (the "Software"),
  # to deal in the Software without restriction, including without limitation
  # the rights to use, copy, modify, merge, publish, distribute, sublicense,
  # and/or sell copies of the Software, and to permit persons to whom the
  # Software is furnished to do so, subject to the following conditions:
  #
  # The above copyright notice and this permission notice shall be included in
  # all copies or substantial portions of the Software.
  #
  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
  # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
  # DEALINGS IN THE SOFTWARE.
  
  import os
  from typing import Any, Callable, Dict, Optional, Sequence
  from llama_index.bridge.pydantic import Field, PrivateAttr
  from llama_index.callbacks import CallbackManager
  from llama_index.constants import DEFAULT_CONTEXT_WINDOW, DEFAULT_NUM_OUTPUTS
  from llama_index.llms.base import (
      ChatMessage,
      ChatResponse,
      CompletionResponse,
      LLMMetadata,
      llm_chat_callback,
      llm_completion_callback,
  )
  from llama_index.llms.custom import CustomLLM
  from llama_index.llms.generic_utils import completion_response_to_chat_response
  from llama_index.llms.generic_utils import (
      messages_to_prompt as generic_messages_to_prompt,
  )
  from transformers import LlamaTokenizer
  import gc
  import json
  import torch
  import numpy as np
  from tensorrt_llm.runtime import ModelConfig, SamplingConfig
  import tensorrt_llm
  from pathlib import Path
  import uuid
  import time
  
  EOS_TOKEN = 2
  PAD_TOKEN = 2
  
  class TrtLlmAPI(CustomLLM):
      model_path: Optional[str] = Field(
          description="The path to the trt engine."
      )
      temperature: float = Field(description="The temperature to use for sampling.")
      max_new_tokens: int = Field(description="The maximum number of tokens to generate.")
      context_window: int = Field(
          description="The maximum number of context tokens for the model."
      )
      messages_to_prompt: Callable = Field(
          description="The function to convert messages to a prompt.", exclude=True
      )
      completion_to_prompt: Callable = Field(
          description="The function to convert a completion to a prompt.", exclude=True
      )
      generate_kwargs: Dict[str, Any] = Field(
          default_factory=dict, description="Kwargs used for generation."
      )
      model_kwargs: Dict[str, Any] = Field(
          default_factory=dict, description="Kwargs used for model initialization."
      )
      verbose: bool = Field(description="Whether to print verbose output.")
  
      _model: Any = PrivateAttr()
      _model_config: Any = PrivateAttr()
      _tokenizer: Any = PrivateAttr()
      _max_new_tokens = PrivateAttr()
      _sampling_config = PrivateAttr()
      _verbose = PrivateAttr()
  
      def __init__(
              self,
              model_path: Optional[str] = None,
              engine_name: Optional[str] = None,
              tokenizer_dir: Optional[str] = None,
              temperature: float = 0.1,
              max_new_tokens: int = DEFAULT_NUM_OUTPUTS,
              context_window: int = DEFAULT_CONTEXT_WINDOW,
              messages_to_prompt: Optional[Callable] = None,
              completion_to_prompt: Optional[Callable] = None,
              callback_manager: Optional[CallbackManager] = None,
              generate_kwargs: Optional[Dict[str, Any]] = None,
              model_kwargs: Optional[Dict[str, Any]] = None,
              verbose: bool = False
      ) -> None:
  
          model_kwargs = model_kwargs or {}
          model_kwargs.update({"n_ctx": context_window, "verbose": verbose})
          self._max_new_tokens = max_new_tokens
          self._verbose = verbose
          # check if model is cached
          if model_path is not None:
              if not os.path.exists(model_path):
                  raise ValueError(
                      "Provided model path does not exist. "
                      "Please check the path or provide a model_url to download."
                  )
              else:
                  engine_dir = model_path
                  engine_dir_path = Path(engine_dir)
                  config_path = engine_dir_path / 'config.json'
  
                  # config function
                  with open(config_path, 'r') as f:
                      config = json.load(f)
                  use_gpt_attention_plugin = config['plugin_config']['gpt_attention_plugin']
                  remove_input_padding = config['plugin_config']['remove_input_padding']
                  tp_size = config['builder_config']['tensor_parallel']
                  pp_size = config['builder_config']['pipeline_parallel']
                  world_size = tp_size * pp_size
                  assert world_size == tensorrt_llm.mpi_world_size(), \
                      f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})'
                  num_heads = config['builder_config']['num_heads'] // tp_size
                  hidden_size = config['builder_config']['hidden_size'] // tp_size
                  vocab_size = config['builder_config']['vocab_size']
                  num_layers = config['builder_config']['num_layers']
                  num_kv_heads = config['builder_config'].get('num_kv_heads', num_heads)
                  paged_kv_cache = config['plugin_config']['paged_kv_cache']
                  if config['builder_config'].get('multi_query_mode', False):
                      tensorrt_llm.logger.warning(
                          "`multi_query_mode` config is deprecated. Please rebuild the engine."
                      )
                      num_kv_heads = 1
                  num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size
  
                  self._model_config = ModelConfig(num_heads=num_heads,
                                                   num_kv_heads=num_kv_heads,
                                                   hidden_size=hidden_size,
                                                   vocab_size=vocab_size,
                                                   num_layers=num_layers,
                                                   gpt_attention_plugin=use_gpt_attention_plugin,
                                                   paged_kv_cache=paged_kv_cache,
                                                   remove_input_padding=remove_input_padding)
  
                  assert pp_size == 1, 'Python runtime does not support pipeline parallelism'
                  world_size = tp_size * pp_size
  
                  runtime_rank = tensorrt_llm.mpi_rank()
                  runtime_mapping = tensorrt_llm.Mapping(world_size,
                                                         runtime_rank,
                                                         tp_size=tp_size,
                                                         pp_size=pp_size)
                  torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node)
                  self._tokenizer = LlamaTokenizer.from_pretrained(tokenizer_dir, legacy=False)
                  self._sampling_config = SamplingConfig(end_id=EOS_TOKEN,
                                                         pad_id=PAD_TOKEN,
                                                         num_beams=1,
                                                         temperature=temperature)
  
                  serialize_path = engine_dir_path / engine_name
                  with open(serialize_path, 'rb') as f:
                      engine_buffer = f.read()
                  decoder = tensorrt_llm.runtime.GenerationSession(self._model_config,
                                                                   engine_buffer,
                                                                   runtime_mapping,
                                                                   debug_mode=False)
                  self._model = decoder
          messages_to_prompt = messages_to_prompt or generic_messages_to_prompt
          completion_to_prompt = completion_to_prompt or (lambda x: x)
  
          generate_kwargs = generate_kwargs or {}
          generate_kwargs.update(
              {"temperature": temperature, "max_tokens": max_new_tokens}
          )
  
          super().__init__(
              model_path=model_path,
              temperature=temperature,
              context_window=context_window,
              max_new_tokens=max_new_tokens,
              messages_to_prompt=messages_to_prompt,
              completion_to_prompt=completion_to_prompt,
              callback_manager=callback_manager,
              generate_kwargs=generate_kwargs,
              model_kwargs=model_kwargs,
              verbose=verbose,
          )
  
      @classmethod
      def class_name(cls) -> str:
          """Get class name."""
          return "TrtLlmAPI"
  
      @property
      def metadata(self) -> LLMMetadata:
          """LLM metadata."""
          return LLMMetadata(
              context_window=self.context_window,
              num_output=self.max_new_tokens,
              model_name=self.model_path,
          )
  
      @llm_chat_callback()
      def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
          prompt = self.messages_to_prompt(messages)
          completion_response = self.complete(prompt, formatted=True, **kwargs)
          return completion_response_to_chat_response(completion_response)
  
      @llm_completion_callback()
      def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
          self.generate_kwargs.update({"stream": False})
  
          is_formatted = kwargs.pop("formatted", False)
          if not is_formatted:
              prompt = self.completion_to_prompt(prompt)
  
          input_text = prompt
          input_ids, input_lengths = self.parse_input(input_text, self._tokenizer,
                                                      EOS_TOKEN,
                                                      self._model_config)
  
          max_input_length = torch.max(input_lengths).item()
          self._model.setup(input_lengths.size(0), max_input_length, self._max_new_tokens, 1) # beam size is set to 1
          if self._verbose:
              start_time = time.time()
  
          output_ids = self._model.decode(input_ids, input_lengths, self._sampling_config)
          torch.cuda.synchronize()
  
          elapsed_time = None
          if self._verbose:
              end_time = time.time()
              elapsed_time = end_time - start_time
  
  
          output_txt, output_token_ids = self.get_output(output_ids,
                                         input_lengths,
                                         self._max_new_tokens,
                                         self._tokenizer)
  
          if self._verbose:
              print(f"Input context length  : {input_ids.shape[1]}")
              print(f"Inference time        : {elapsed_time:.2f} seconds")
              print(f"Output context length : {len(output_token_ids)} ")
              print(f"Inference token/sec   : {(len(output_token_ids) / elapsed_time):2f}")
  
          # call garbage collected after inference
          torch.cuda.empty_cache()
          gc.collect()
  
          return CompletionResponse(text=output_txt, raw=self.generate_completion_dict(output_txt))
  
      def parse_input(self, input_text: str, tokenizer, end_id: int,
                      remove_input_padding: bool):
          input_tokens = []
  
          input_tokens.append(
              tokenizer.encode(input_text, add_special_tokens=False))
  
          input_lengths = torch.tensor([len(x) for x in input_tokens],
                                       dtype=torch.int32,
                                       device='cuda')
          if remove_input_padding:
              input_ids = np.concatenate(input_tokens)
              input_ids = torch.tensor(input_ids, dtype=torch.int32,
                                       device='cuda').unsqueeze(0)
          else:
              input_ids = torch.nested.to_padded_tensor(
                  torch.nested.nested_tensor(input_tokens, dtype=torch.int32),
                  end_id).cuda()
  
          return input_ids, input_lengths
  
      def remove_extra_eos_ids(self, outputs):
          outputs.reverse()
          while outputs and outputs[0] == 2:
              outputs.pop(0)
          outputs.reverse()
          outputs.append(2)
          return outputs
  
      def get_output(self, output_ids, input_lengths, max_output_len, tokenizer):
          num_beams = output_ids.size(1)
          output_text = ""
          outputs = None
          for b in range(input_lengths.size(0)):
              for beam in range(num_beams):
                  output_begin = input_lengths[b]
                  output_end = input_lengths[b] + max_output_len
                  outputs = output_ids[b][beam][output_begin:output_end].tolist()
                  outputs = self.remove_extra_eos_ids(outputs)
                  output_text = tokenizer.decode(outputs)
  
          return output_text, outputs
  
      def generate_completion_dict(self, text_str):
          """
          Generate a dictionary for text completion details.
          Returns:
          dict: A dictionary containing completion details.
          """
          completion_id: str = f"cmpl-{str(uuid.uuid4())}"
          created: int = int(time.time())
          model_name: str = self._model if self._model is not None else self.model_path
          return {
              "id": completion_id,
              "object": "text_completion",
              "created": created,
              "model": model_name,
              "choices": [
                  {
                      "text": text_str,
                      "index": 0,
                      "logprobs": None,
                      "finish_reason": 'stop'
                  }
              ],
              "usage": {
                  "prompt_tokens": None,
                  "completion_tokens": None,
                  "total_tokens": None
              }
          }
  
      @llm_completion_callback()
      def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
          pass
  
apis: llama_index.llms.base.LLMMetadata llama_index.bridge.pydantic.Field llama_index.llms.base.llm_completion_callback llama_index.bridge.pydantic.PrivateAttr llama_index.llms.base.llm_chat_callback llama_index.llms.generic_utils.completion_response_to_chat_response
extract_api: [((2151, 2199), 'llama_index.bridge.pydantic.Field', 'Field', ([], {'description': '"""The path to the trt engine."""'}), "(description='The path to the trt engine.')\n", (2156, 2199), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((2239, 2296), 'llama_index.bridge.pydantic.Field', 'Field', ([], {'description': '"""The temperature to use for sampling."""'}), "(description='The temperature to use for sampling.')\n", (2244, 2296), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((2323, 2385), 'llama_index.bridge.pydantic.Field', 'Field', ([], {'description': '"""The maximum number of tokens to generate."""'}), "(description='The maximum number of tokens to generate.')\n", (2328, 2385), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((2412, 2484), 'llama_index.bridge.pydantic.Field', 'Field', ([], {'description': '"""The maximum number of context tokens for the model."""'}), "(description='The maximum number of context tokens for the model.')\n", (2417, 2484), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((2534, 2619), 'llama_index.bridge.pydantic.Field', 'Field', ([], {'description': '"""The function to convert messages to a prompt."""', 'exclude': '(True)'}), "(description='The function to convert messages to a prompt.', exclude=True\n    )\n", (2539, 2619), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((2666, 2754), 'llama_index.bridge.pydantic.Field', 'Field', ([], {'description': '"""The function to convert a completion to a prompt."""', 'exclude': '(True)'}), "(description='The function to convert a completion to a prompt.',\n    exclude=True)\n", (2671, 2754), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((2803, 2873), 'llama_index.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'dict', 'description': '"""Kwargs used for generation."""'}), "(default_factory=dict, description='Kwargs used for generation.')\n", (2808, 2873), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((2923, 3008), 'llama_index.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'dict', 'description': '"""Kwargs used for model initialization."""'}), "(default_factory=dict, description='Kwargs used for model initialization.'\n    )\n", (2928, 3008), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((3038, 3091), 'llama_index.bridge.pydantic.Field', 'Field', ([], {'description': '"""Whether to print verbose output."""'}), "(description='Whether to print verbose output.')\n", (3043, 3091), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((3111, 3124), 'llama_index.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (3122, 3124), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((3150, 3163), 'llama_index.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (3161, 3163), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((3186, 3199), 'llama_index.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (3197, 3199), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((3222, 3235), 'llama_index.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (3233, 3235), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((3259, 3272), 'llama_index.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (3270, 3272), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((3288, 3301), 'llama_index.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (3299, 3301), False, 'from llama_index.bridge.pydantic import Field, PrivateAttr\n'), ((9389, 9408), 'llama_index.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (9406, 9408), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, CompletionResponse, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((9701, 9726), 'llama_index.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (9724, 9726), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, CompletionResponse, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((14134, 14159), 'llama_index.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (14157, 14159), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, CompletionResponse, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((9228, 9340), 'llama_index.llms.base.LLMMetadata', 'LLMMetadata', ([], {'context_window': 'self.context_window', 'num_output': 'self.max_new_tokens', 'model_name': 'self.model_path'}), '(context_window=self.context_window, num_output=self.\n    max_new_tokens, model_name=self.model_path)\n', (9239, 9340), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, CompletionResponse, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((9637, 9694), 'llama_index.llms.generic_utils.completion_response_to_chat_response', 'completion_response_to_chat_response', (['completion_response'], {}), '(completion_response)\n', (9673, 9694), False, 'from llama_index.llms.generic_utils import completion_response_to_chat_response\n'), ((10577, 10601), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ([], {}), '()\n', (10599, 10601), False, 'import torch\n'), ((11367, 11391), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ([], {}), '()\n', (11389, 11391), False, 'import torch\n'), ((11400, 11412), 'gc.collect', 'gc.collect', ([], {}), '()\n', (11410, 11412), False, 'import gc\n'), ((10467, 10478), 'time.time', 'time.time', ([], {}), '()\n', (10476, 10478), False, 'import time\n'), ((10680, 10691), 'time.time', 'time.time', ([], {}), '()\n', (10689, 10691), False, 'import time\n'), ((11988, 12016), 'numpy.concatenate', 'np.concatenate', (['input_tokens'], {}), '(input_tokens)\n', (12002, 12016), True, 'import numpy as np\n'), ((13479, 13490), 'time.time', 'time.time', ([], {}), '()\n', (13488, 13490), False, 'import time\n'), ((4271, 4297), 'os.path.exists', 'os.path.exists', (['model_path'], {}), '(model_path)\n', (4285, 4297), False, 'import os\n'), ((4582, 4598), 'pathlib.Path', 'Path', (['engine_dir'], {}), '(engine_dir)\n', (4586, 4598), False, 'from pathlib import Path\n'), ((6180, 6445), 'tensorrt_llm.runtime.ModelConfig', 'ModelConfig', ([], {'num_heads': 'num_heads', 'num_kv_heads': 'num_kv_heads', 'hidden_size': 'hidden_size', 'vocab_size': 'vocab_size', 'num_layers': 'num_layers', 'gpt_attention_plugin': 'use_gpt_attention_plugin', 'paged_kv_cache': 'paged_kv_cache', 'remove_input_padding': 'remove_input_padding'}), '(num_heads=num_heads, num_kv_heads=num_kv_heads, hidden_size=\n    hidden_size, vocab_size=vocab_size, num_layers=num_layers,\n    gpt_attention_plugin=use_gpt_attention_plugin, paged_kv_cache=\n    paged_kv_cache, remove_input_padding=remove_input_padding)\n', (6191, 6445), False, 'from tensorrt_llm.runtime import ModelConfig, SamplingConfig\n'), ((6947, 6970), 'tensorrt_llm.mpi_rank', 'tensorrt_llm.mpi_rank', ([], {}), '()\n', (6968, 6970), False, 'import tensorrt_llm\n'), ((7005, 7090), 'tensorrt_llm.Mapping', 'tensorrt_llm.Mapping', (['world_size', 'runtime_rank'], {'tp_size': 'tp_size', 'pp_size': 'pp_size'}), '(world_size, runtime_rank, tp_size=tp_size, pp_size=pp_size\n    )\n', (7025, 7090), False, 'import tensorrt_llm\n'), ((7267, 7334), 'torch.cuda.set_device', 'torch.cuda.set_device', (['(runtime_rank % runtime_mapping.gpus_per_node)'], {}), '(runtime_rank % runtime_mapping.gpus_per_node)\n', (7288, 7334), False, 'import torch\n'), ((7369, 7428), 'transformers.LlamaTokenizer.from_pretrained', 'LlamaTokenizer.from_pretrained', (['tokenizer_dir'], {'legacy': '(False)'}), '(tokenizer_dir, legacy=False)\n', (7399, 7428), False, 'from transformers import LlamaTokenizer\n'), ((7469, 7562), 'tensorrt_llm.runtime.SamplingConfig', 'SamplingConfig', ([], {'end_id': 'EOS_TOKEN', 'pad_id': 'PAD_TOKEN', 'num_beams': '(1)', 'temperature': 'temperature'}), '(end_id=EOS_TOKEN, pad_id=PAD_TOKEN, num_beams=1, temperature\n    =temperature)\n', (7483, 7562), False, 'from tensorrt_llm.runtime import ModelConfig, SamplingConfig\n'), ((7912, 8024), 'tensorrt_llm.runtime.GenerationSession', 'tensorrt_llm.runtime.GenerationSession', (['self._model_config', 'engine_buffer', 'runtime_mapping'], {'debug_mode': '(False)'}), '(self._model_config, engine_buffer,\n    runtime_mapping, debug_mode=False)\n', (7950, 8024), False, 'import tensorrt_llm\n'), ((10268, 10292), 'torch.max', 'torch.max', (['input_lengths'], {}), '(input_lengths)\n', (10277, 10292), False, 'import torch\n'), ((4775, 4787), 'json.load', 'json.load', (['f'], {}), '(f)\n', (4784, 4787), False, 'import json\n'), ((5192, 5221), 'tensorrt_llm.mpi_world_size', 'tensorrt_llm.mpi_world_size', ([], {}), '()\n', (5219, 5221), False, 'import tensorrt_llm\n'), ((5889, 5992), 'tensorrt_llm.logger.warning', 'tensorrt_llm.logger.warning', (['"""`multi_query_mode` config is deprecated. Please rebuild the engine."""'], {}), "(\n    '`multi_query_mode` config is deprecated. Please rebuild the engine.')\n", (5916, 5992), False, 'import tensorrt_llm\n'), ((12041, 12098), 'torch.tensor', 'torch.tensor', (['input_ids'], {'dtype': 'torch.int32', 'device': '"""cuda"""'}), "(input_ids, dtype=torch.int32, device='cuda')\n", (12053, 12098), False, 'import torch\n'), ((13436, 13448), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (13446, 13448), False, 'import uuid\n'), ((5304, 5333), 'tensorrt_llm.mpi_world_size', 'tensorrt_llm.mpi_world_size', ([], {}), '()\n', (5331, 5333), False, 'import tensorrt_llm\n'), ((12234, 12293), 'torch.nested.nested_tensor', 'torch.nested.nested_tensor', (['input_tokens'], {'dtype': 'torch.int32'}), '(input_tokens, dtype=torch.int32)\n', (12260, 12293), False, 'import torch\n')]