class_docstr: Model wrapper around HuggingFace AutoModel models.
  
  Args:
      path: The path to a HuggingFace model.
      model_max_length: The maximum sequence length of the model.
      tokenizer_path: The path to the tokenizer.
      tokenizer_kwargs: Keyword arguments for the tokenizer.
      peft_path: The name or path to the HuggingFace's PEFT model.
      model_kwargs: Keyword arguments for the model.
      prompt_template: The model's prompt template.
      batch_size: Batch size for inference.
      logger: Logger for the model.
      shard_config: Shard config for tensor parallel.
class_docstr_tok: ["Model","wrapper","around","HuggingFace","AutoModel","models",".","Args",":","path",":","The","path","to","a","HuggingFace","model",".","model_max_length",":","The","maximum","sequence","length","of","the","model",".","tokenizer_path",":","The","path","to","the","tokenizer",".","tokenizer_kwargs",":","Keyword","arguments","for","the","tokenizer",".","peft_path",":","The","name","or","path","to","the","HuggingFace","'s","PEFT","model",".","model_kwargs",":","Keyword","arguments","for","the","model",".","prompt_template",":","The","model","'s","prompt","template",".","batch_size",":","Batch","size","for","inference",".","logger",":","Logger","for","the","model",".","shard_config",":","Shard","config","for","tensor","parallel","."]
docstr: Calculate loss only on target tokens.
  
  Args:
      batch: A batch of prompt without target answer.
      batch_target: A batch of target answer. Sometimes one question can have multiple target answers.
  
  Returns:
      Loss.
docstr_tok: ["Calculate","loss","only","on","target","tokens",".","Args",":","batch",":","A","batch","of","prompt","without","target","answer",".","batch_target",":","A","batch","of","target","answer",".","Sometimes","one","question","can","have","multiple","target","answers",".","Returns",":","Loss","."]
code: def get_loss(
      self,
      batch_prompt: List[str],
      batch_target: List[List[str]],
      pretrain: bool,
  ) -> List[List[float]]:
      """
      Calculate loss only on target tokens.
  
      Args:
          batch: A batch of prompt without target answer.
          batch_target: A batch of target answer. Sometimes one question can have multiple target answers.
  
      Returns:
          Loss.
  
      """
  
      # We set max_new_tokens in self._get_truncated_prompts to 0 because we only need logits to calculate loss.
      # We don't need to generate new tokens.
      # Target answer's length is usually << model_max_length, but we still call it in case.
      # We don't call self._get_truncated_prompts for batch_prompt because we need target answer's length first to reserve some space for target answer's tokens.
      if not pretrain:
          batch_target = [
              self._get_truncated_prompts(prompt_target, 0)
              for prompt_target in batch_target
          ]
  
      # Get the number of target answers for different questions
      batch_target_nums = [
          len(prompt_target) for prompt_target in batch_target
      ]
  
      (
          input_ids_list,
          labels_list,
          bytes_list,
      ) = self._get_input_ids_and_labels(
          batch_prompt, batch_target, pretrain
      )
  
      # Because of multiple target answers, the final batch size may be greater than self.batch_size.
      # We will generate new batches.
      losses = []
      target_token_nums = []
  
      batched_input_ids = [
          input_ids_list[i : i + self.batch_size]
          for i in range(0, len(input_ids_list), self.batch_size)
      ]
      batched_labels = [
          labels_list[i : i + self.batch_size]
          for i in range(0, len(labels_list), self.batch_size)
      ]
  
      for batch_input_ids, batch_labels in zip(
          batched_input_ids, batched_labels
      ):
          (
              losses_per_batch,
              target_token_num_per_batch,
          ) = self._calculate_loss(batch_input_ids, batch_labels)
          losses.extend(losses_per_batch)
          target_token_nums.extend(target_token_num_per_batch)
  
      start_indice = 0
      losses_per_sample = []
  
      target_token_nums_per_sample = []
      bytes_nums_per_sample = []
      for length in batch_target_nums:
          losses_per_sample.append(
              losses[start_indice : start_indice + length]
          )
          target_token_nums_per_sample.append(
              target_token_nums[start_indice : start_indice + length]
          )
  
          if bytes_list:
              bytes_nums_per_sample.append(
                  bytes_list[start_indice : start_indice + length]
              )
  
          start_indice += length
  
      if bytes_list:
          return (
              losses_per_sample,
              target_token_nums_per_sample,
              bytes_nums_per_sample,
          )
  
      return losses_per_sample, target_token_nums_per_sample, None
  
imports: import copy
  import math
  from typing import Any, Dict, List, Optional, Tuple
  import numpy
  import torch
  from colossal_eval.utils import Conversation, get_batch_prompt, is_rank_0
  from peft import PeftModel
  from tqdm import tqdm
  from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer
  from colossalai.logging import DistributedLogger
  from colossalai.shardformer import ShardConfig, ShardFormer
  from .base import BaseModel
label_desc: Use image node_id 11 for calling the HuggingFaceModel obj's underlying member method code with example usage: obj.get_loss(batch_prompt, batch_target, pretrain) and returns: losses_per_sample,  target_token_nums_per_sample,  None, losses_per_sample,  target_token_nums_per_sample,  bytes_nums_per_sample