QC_Specialized_Model_Training_Main/experiments/inference_formal.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.

import os
import shutil
import sys
import time
import fire
import torch
from tqdm import tqdm

from accelerate.utils import is_xpu_available
from llama_recipes.inference.model_utils import load_model, load_peft_model

from llama_recipes.inference.safety_utils import AgentType, get_safety_checker
from transformers import AutoTokenizer

from llama_recipes.data.concatenator import ConcatDataset

from llama_recipes.utils.dataset_utils import get_enote_dataset

from llama_recipes.configs import (
    fsdp_config as FSDP_CONFIG,
    quantization_config as QUANTIZATION_CONFIG,
    train_config as TRAIN_CONFIG,
)

from llama_recipes.utils.config_utils import (
    check_fsdp_config,
    generate_dataset_config,
    generate_peft_config,
    get_dataloader_kwargs,
    update_config,
)


def create_clean_dir(path):
    """
    Create a clean directory. If the directory exists, remove it first.
    :param path: Path of the directory to create.
    """
    # Remove the directory if it exists
    if os.path.exists(path):
        shutil.rmtree(path)
    # Create the directory
    os.makedirs(path)


def main(
    model_name,
    peft_model: str = None,
    quantization: str = None,  # Options: 4bit, 8bit
    max_new_tokens=1000,  # The maximum numbers of tokens to generate
    prompt_file: str = None,
    seed: int = 42,  # seed value for reproducibility
    do_sample: bool = True,  # Whether or not to use sampling ; use greedy decoding otherwise.
    min_length: int = None,  # The minimum length of the sequence to be generated, input prompt + min_new_tokens
    use_cache: bool = True,
    # [optional] Whether or not the model should use the past last key/values attentions Whether or not the model should use the past last key/values attentions (if applicable to the model) to speed up decoding.
    top_p: float = 1.0,
    # [optional] If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation.
    temperature: float = 1.0,  # [optional] The value used to modulate the next token probabilities.
    top_k: int = 50,  # [optional] The number of highest probability vocabulary tokens to keep for top-k-filtering.
    repetition_penalty: float = 1.0,  # The parameter for repetition penalty. 1.0 means no penalty.
    length_penalty: int = 1,
    # [optional] Exponential penalty to the length that is used with beam-based generation.
    enable_azure_content_safety: bool = False,  # Enable safety check with Azure content safety api
    enable_sensitive_topics: bool = False,  # Enable check for sensitive topics using AuditNLG APIs
    enable_salesforce_content_safety: bool = True,  # Enable safety check with Salesforce safety flan t5
    enable_llamaguard_content_safety: bool = False,
    max_padding_length: int = None,  # the max padding length to be used with tokenizer padding the prompts.
    use_fast_kernels: bool = False,
    # Enable using SDPA from PyTroch Accelerated Transformers, make use Flash Attention and Xformer memory-efficient kernels
    share_gradio: bool = False,  # Enable endpoint creation for gradio.live
    lang_pairs: str = None,
    output_dir: str = None,
    **kwargs,
):
    # Set the seeds for reproducibility
    if is_xpu_available():
        torch.xpu.manual_seed(seed)
    else:
        torch.cuda.manual_seed(seed)
    torch.manual_seed(seed)

    # Update the configuration for the training and sharding process
    test_config, fsdp_config = TRAIN_CONFIG(), FSDP_CONFIG()
    update_config((test_config, fsdp_config), **kwargs)
    # dataset_config = generate_dataset_config(test_config, kwargs)

    model = load_model(model_name, quantization, use_fast_kernels, **kwargs)

    if test_config.preload_peft_dir is not None:
        # merge peft into backbone, may not 100% aligned
        print("Load and merge peft...")
        model = load_peft_model(model, test_config.preload_peft_dir)
        model = model.merge_and_unload()

    if peft_model:
        model = load_peft_model(model, peft_model)

    model.eval()

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = 'left'
    model.generation_config.pad_token_id = tokenizer.pad_token_id

    # TODO, batch inference
    def inference_new(
            dataloader,
            temperature,
            top_p,
            top_k,
            max_new_tokens,
            config,
            pbar,
            **kwargs,
    ):
        output = []
        for step, batch in enumerate(dataloader):
            # TODO, dirty
            batch.pop('labels')
            if is_xpu_available():
                batch = {k: v.to("xpu") for k, v in batch.items()}
            else:
                batch = {k: v.to("cuda") for k, v in batch.items()}

            with torch.no_grad():
                batch_output = model.generate(
                    **batch,
                    max_new_tokens=max_new_tokens,
                    do_sample=do_sample,
                    top_p=top_p,
                    temperature=temperature,
                    min_length=min_length,
                    use_cache=use_cache,
                    top_k=top_k,
                    repetition_penalty=repetition_penalty,
                    length_penalty=length_penalty,
                    num_beams=config.beam_size,
                    **kwargs,
                )

                # prompt_len = batch['attention_mask'].sum(-1)

                batch_len = batch['input_ids'].shape[-1]
                batch_output = batch_output[:, batch_len:]
                batch_output = [tokenizer.decode(output, skip_special_tokens=True) for output in batch_output]

                # replace \n with \t to read when hallucinating
                batch_output = [sent.replace("\n", "\t").strip() for sent in batch_output]

                output += batch_output
            pbar.update(1)
        return output

    # TODO, inference for each dataset
    output = {}

    rule_names = test_config.rule_names

    for rule_name in rule_names:
        # Get test data
        print("Processing {} ...".format(rule_name), flush=True)

        dataset_test = get_enote_dataset(
            tokenizer,
            test_config.dataset,
            mode="infer",
            split="valid",
            rule_names=rule_names
        )
        print(f"--> Test Set Length = {len(dataset_test)}", flush=True)

        test_dl_kwargs = get_dataloader_kwargs(
            test_config, dataset_test, tokenizer, "infer"
        )

        # Create DataLoaders for inference
        test_dataloader = torch.utils.data.DataLoader(
            dataset_test,
            num_workers=0,
            pin_memory=True,
            shuffle=False,
            **test_dl_kwargs,
        )
        print(f"--> Num of Testing Set Batches loaded = {len(test_dataloader)}", flush=True)

        start = time.perf_counter()
        total_length = len(test_dataloader)
        pbar = tqdm(colour="blue", desc=f"Inference", total=total_length, dynamic_ncols=True)
        results = inference_new(test_dataloader, temperature, top_p, top_k, max_new_tokens, test_config, pbar=pbar)
        pbar.close()
        e2e_inference_time = (time.perf_counter() - start) * 1000
        print(f"the inference time is {e2e_inference_time} ms", flush=True)
        output[rule_name] = results

        # dump results
        create_clean_dir(os.path.join(output_dir, rule_name))
        output_file = os.path.join(output_dir, rule_name, "hyp.{}".format(rule_name))
        with open(output_file, 'w') as fout:
            for line in results:
                fout.write(line.strip() + "\n")


if __name__ == "__main__":
    fire.Fire(main)