unitorch.models.llama

LlamaProcessor

Bases: HfTextClassificationProcessor, HfTextGenerationProcessor

Source code in src/unitorch/models/llama/processing.py

def __init__(
    self,
    tokenizer_file: str,
    tokenizer_config: Optional[str] = None,
    special_tokens_map: Optional[str] = None,
    chat_template: Optional[str] = None,
    max_seq_length: Optional[int] = 128,
    max_gen_seq_length: Optional[int] = 48,
):
    tokenizer_config = read_json_file(tokenizer_config) if tokenizer_config else {}
    special_tokens_map = (
        read_json_file(special_tokens_map) if special_tokens_map else {}
    )
    added_tokens_decoder = tokenizer_config.pop("added_tokens_decoder", {})
    tokenizer_config = {
        k: (
            get_added_token(v)
            if isinstance(v, dict) and v.get("__type") == "AddedToken"
            else v
        )
        for k, v in tokenizer_config.items()
    }

    tokenizer = LlamaTokenizerFast(
        tokenizer_file=tokenizer_file,
        **tokenizer_config,
    )

    for idx, spec in added_tokens_decoder.items():
        token = spec["content"]
        tokenizer.added_tokens_decoder[idx] = get_added_token(spec)
        tokenizer.added_tokens_encoder[token] = idx

    special_tokens = {}
    for name, spec in special_tokens_map.items():
        special_tokens[name] = get_added_token(spec)
    tokenizer.add_special_tokens(special_tokens)

    if chat_template:
        tokenizer.chat_template = read_json_file(chat_template)["chat_template"]
    tokenizer.cls_token = tokenizer.bos_token
    tokenizer.sep_token = tokenizer.eos_token
    tokenizer.pad_token = (
        tokenizer.unk_token if tokenizer.pad_token is None else tokenizer.pad_token
    )
    tokenizer.cls_token_id = tokenizer.bos_token_id
    tokenizer.sep_token_id = tokenizer.eos_token_id
    tokenizer.pad_token_id = (
        tokenizer.unk_token_id
        if tokenizer.pad_token_id is None
        else tokenizer.pad_token_id
    )
    HfTextClassificationProcessor.__init__(
        self,
        tokenizer=tokenizer,
        max_seq_length=max_seq_length,
    )
    HfTextGenerationProcessor.__init__(
        self,
        tokenizer=tokenizer,
        max_seq_length=max_seq_length,
        max_gen_seq_length=max_gen_seq_length,
    )

chat_template

chat_template(messages: List[Dict[str, Any]])

Source code in src/unitorch/models/llama/processing.py

def chat_template(
    self,
    messages: List[Dict[str, Any]],
):
    text = self.tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    return text

classification

classification(
    text: str,
    text_pair: Optional[str] = None,
    max_seq_length: Optional[int] = None,
)

Source code in src/unitorch/models/llama/processing.py

def classification(
    self,
    text: str,
    text_pair: Optional[str] = None,
    max_seq_length: Optional[int] = None,
):
    max_seq_length = pop_value(
        max_seq_length,
        self.max_seq_length,
    )

    tokens = self.tokenizer.tokenize(str(text))
    if text_pair is None:
        tokens = tokens[:max_seq_length]
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
    else:
        tokens_pair = self.tokenizer.tokenize(str(text_pair))
        truncate_sequence_pair(tokens, tokens_pair, max_seq_length)
        tokens = tokens + tokens_pair
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)

    padding = [0] * (max_seq_length - len(input_ids))
    attention_mask = [0] * len(padding) + [1] * len(input_ids)
    input_ids = len(padding) * [self.pad_token_id] + input_ids

    assert len(input_ids) == max_seq_length
    assert len(attention_mask) == max_seq_length
    return GenericOutputs(
        input_ids=torch.tensor(input_ids, dtype=torch.long),
        attention_mask=torch.tensor(attention_mask, dtype=torch.long),
    )

generation_inputs

generation_inputs(
    text: str, max_seq_length: Optional[int] = None
)

Source code in src/unitorch/models/llama/processing.py

def generation_inputs(
    self,
    text: str,
    max_seq_length: Optional[int] = None,
):
    max_seq_length = pop_value(
        max_seq_length,
        self.max_seq_length,
    )
    tokens = [self.bos_token] + self.tokenizer.tokenize(str(text))[
        1 - max_seq_length :
    ]
    padding = [self.pad_token] * (max_seq_length - len(tokens))
    tokens = padding + tokens
    input_ids = self.tokenizer.convert_tokens_to_ids(tokens)

    assert len(input_ids) == max_seq_length
    return GenericOutputs(
        input_ids=torch.tensor(input_ids, dtype=torch.long),
    )

generation_labels

generation_labels(
    text: str, max_gen_seq_length: Optional[int] = None
)

Source code in src/unitorch/models/llama/processing.py

def generation_labels(
    self,
    text: str,
    max_gen_seq_length: Optional[int] = None,
):
    max_gen_seq_length = pop_value(
        max_gen_seq_length,
        self.max_gen_seq_length,
    )
    tokens = self.tokenizer.tokenize(str(text))[: max_gen_seq_length - 1] + [
        self.eos_token
    ]
    padding = [self.pad_token] * (max_gen_seq_length - len(tokens))
    input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
    attention_mask = [1] * len(input_ids)

    padding = [0] * (max_gen_seq_length - len(input_ids))
    input_ids += [self.pad_token_id] * len(padding)
    attention_mask += padding

    assert len(input_ids) == max_gen_seq_length
    assert len(attention_mask) == max_gen_seq_length
    return GenericOutputs(
        input_ids=torch.tensor(input_ids, dtype=torch.long),
        attention_mask=torch.tensor(attention_mask, dtype=torch.long),
    )

generation

generation(
    text: str,
    text_pair: str,
    max_seq_length: Optional[int] = None,
    max_gen_seq_length: Optional[int] = None,
)

Source code in src/unitorch/models/llama/processing.py

def generation(
    self,
    text: str,
    text_pair: str,
    max_seq_length: Optional[int] = None,
    max_gen_seq_length: Optional[int] = None,
):
    max_seq_length = pop_value(
        max_seq_length,
        self.max_seq_length,
    )
    max_gen_seq_length = pop_value(
        max_gen_seq_length,
        self.max_gen_seq_length,
    )

    tokens = [self.bos_token] + self.tokenizer.tokenize(str(text))[
        1 - max_seq_length :
    ]
    tokens_pair = self.tokenizer.tokenize(str(text_pair))[
        : max_gen_seq_length - 1
    ] + [self.eos_token]
    padding_a = [self.pad_token] * (max_seq_length - len(tokens))
    padding_b = [self.pad_token] * (max_gen_seq_length - len(tokens_pair))
    attention_mask = (
        [0] * len(padding_a)
        + [1] * (len(tokens) + len(tokens_pair))
        + [0] * len(padding_b)
    )
    _tokens = padding_a + tokens + tokens_pair + padding_b
    input_ids = self.tokenizer.convert_tokens_to_ids(_tokens)

    tokens_label = tokens_pair + [self.pad_token] * (
        max_gen_seq_length - len(tokens_pair) + 1
    )
    input_ids_label = self.tokenizer.convert_tokens_to_ids(tokens_label)
    input_ids_label = [0] * (max_seq_length - 1) + input_ids_label
    attention_mask_label = [1] * len(tokens_pair) + [0] * (
        max_gen_seq_length - len(tokens_pair) + 1
    )
    attention_mask_label = [0] * (max_seq_length - 1) + attention_mask_label

    return GenericOutputs(
        input_ids=torch.tensor(input_ids, dtype=torch.long),
        attention_mask=torch.tensor(attention_mask, dtype=torch.long),
        input_ids_label=torch.tensor(input_ids_label, dtype=torch.long),
        attention_mask_label=torch.tensor(attention_mask_label, dtype=torch.long),
    )

messages_generation

messages_generation(
    messages: List[Dict[str, Any]],
    max_seq_length: Optional[int] = None,
) -> GenericOutputs

Source code in src/unitorch/models/llama/processing.py

def messages_generation(
    self,
    messages: List[Dict[str, Any]],
    max_seq_length: Optional[int] = None,
) -> GenericOutputs:
    while messages and messages[-1]["role"] != "assistant":
        messages.pop()

    text = self.chat_template(messages[:-1])
    text_pair = self.chat_template(messages[-1:])
    outputs = self.generation(
        text=text,
        text_pair=text_pair,
        max_seq_length=max_seq_length,
    )
    return GenericOutputs(
        input_ids=outputs.input_ids,
        attention_mask=outputs.attention_mask,
        input_ids_label=outputs.input_ids_label,
        attention_mask_label=outputs.attention_mask_label,
    )

LlamaForClassification

Bases: GenericModel, PeftWeightLoaderMixin

Source code in src/unitorch/models/llama/modeling.py

def __init__(
    self,
    config_path: str,
    num_classes: Optional[int] = 1,
    hidden_dropout_prob: Optional[float] = 0.1,
    gradient_checkpointing: Optional[bool] = False,
):
    super().__init__()
    self.config = LlamaConfig.from_json_file(config_path)
    self.config.gradient_checkpointing = gradient_checkpointing
    self.model = LlamaModel(self.config)
    self.dropout = nn.Dropout(hidden_dropout_prob)
    self.classifier = nn.Linear(self.config.hidden_size, num_classes)
    self.init_weights()

config instance-attribute

config = from_json_file(config_path)

model instance-attribute

model = LlamaModel(config)

dropout instance-attribute

dropout = Dropout(hidden_dropout_prob)

classifier instance-attribute

classifier = Linear(hidden_size, num_classes)

forward

forward(
    input_ids: Tensor,
    attention_mask: Optional[Tensor] = None,
    position_ids: Optional[Tensor] = None,
)

Source code in src/unitorch/models/llama/modeling.py

def forward(
    self,
    input_ids: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.Tensor] = None,
):
    outputs = self.model(
        input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
    )[0]
    pooled_output = outputs[:, -1]
    pooled_output = self.dropout(pooled_output)
    logits = self.classifier(pooled_output)
    return logits

LlamaForGeneration

Bases: GenericModel, PeftWeightLoaderMixin

Source code in src/unitorch/models/llama/modeling.py

def __init__(
    self,
    config_path: str,
    gradient_checkpointing: Optional[bool] = False,
):
    super().__init__()
    self.config = LlamaConfig.from_json_file(config_path)
    self.config.gradient_checkpointing = gradient_checkpointing
    self.base_model = LlamaForCausalLM(self.config)
    self.init_weights()

prefix_keys_in_state_dict class-attribute instance-attribute

prefix_keys_in_state_dict = {
    "^model.": "base_model.",
    "^lm_head.": "base_model.",
}

config instance-attribute

config = from_json_file(config_path)

base_model instance-attribute

base_model = LlamaForCausalLM(config)

forward

forward(
    input_ids: Tensor,
    attention_mask: Optional[Tensor] = None,
    position_ids: Optional[Tensor] = None,
)

Source code in src/unitorch/models/llama/modeling.py

def forward(
    self,
    input_ids: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.Tensor] = None,
):
    outputs = self.base_model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        return_dict=True,
    )
    logits = outputs.logits
    return logits

generate

generate(
    input_ids: Tensor,
    num_beams: Optional[int] = 5,
    decoder_start_token_id: Optional[int] = 1,
    decoder_end_token_id: Optional[
        Union[int, List[int]]
    ] = 2,
    decoder_pad_token_id: Optional[int] = 1,
    num_return_sequences: Optional[int] = 1,
    min_gen_seq_length: Optional[int] = 0,
    max_gen_seq_length: Optional[int] = 48,
    repetition_penalty: Optional[float] = 1.0,
    no_repeat_ngram_size: Optional[int] = 0,
    early_stopping: Optional[bool] = True,
    length_penalty: Optional[float] = 1.0,
    num_beam_groups: Optional[int] = 1,
    diversity_penalty: Optional[float] = 0.0,
    do_sample: Optional[bool] = False,
    temperature: Optional[float] = 1.0,
    top_k: Optional[int] = 50,
    top_p: Optional[float] = 1.0,
)

Source code in src/unitorch/models/llama/modeling.py

@torch.no_grad()
def generate(
    self,
    input_ids: torch.Tensor,
    num_beams: Optional[int] = 5,
    decoder_start_token_id: Optional[int] = 1,
    decoder_end_token_id: Optional[Union[int, List[int]]] = 2,
    decoder_pad_token_id: Optional[int] = 1,
    num_return_sequences: Optional[int] = 1,
    min_gen_seq_length: Optional[int] = 0,
    max_gen_seq_length: Optional[int] = 48,
    repetition_penalty: Optional[float] = 1.0,
    no_repeat_ngram_size: Optional[int] = 0,
    early_stopping: Optional[bool] = True,
    length_penalty: Optional[float] = 1.0,
    num_beam_groups: Optional[int] = 1,
    diversity_penalty: Optional[float] = 0.0,
    do_sample: Optional[bool] = False,
    temperature: Optional[float] = 1.0,
    top_k: Optional[int] = 50,
    top_p: Optional[float] = 1.0,
):
    input_seq_length = input_ids.size(1)
    outputs = self.base_model.generate(
        input_ids,
        max_length=max_gen_seq_length + input_seq_length,
        min_length=min_gen_seq_length + input_seq_length,
        num_beams=num_beams,
        do_sample=do_sample,
        no_repeat_ngram_size=no_repeat_ngram_size,
        early_stopping=early_stopping,
        length_penalty=length_penalty,
        repetition_penalty=repetition_penalty,
        num_return_sequences=num_return_sequences,
        bos_token_id=decoder_start_token_id,
        eos_token_id=decoder_end_token_id,
        pad_token_id=decoder_pad_token_id,
        num_beam_groups=num_beam_groups,
        diversity_penalty=diversity_penalty,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        return_dict_in_generate=True,
        output_scores=True,
    )

    sequences = outputs.sequences.reshape(
        -1, num_return_sequences, outputs.sequences.size(-1)
    )
    outputs.sequences = (
        torch.zeros(sequences.size(0), num_return_sequences, max_gen_seq_length).to(
            device=sequences.device
        )
        + decoder_start_token_id
    )
    outputs.sequences[:, :, : sequences.size(-1) - input_seq_length].copy_(
        sequences[:, :, input_seq_length : sequences.size(-1)]
    )

    if num_return_sequences == 1:
        outputs.sequences = outputs.sequences.reshape(-1, max_gen_seq_length)

    return GenericOutputs(
        sequences=outputs.sequences.long(),
        sequences_scores=outputs.sequences_scores,
    )