# This script is used to generate predictions using the fine-tuned adapter models

# This script is modified from the original script provided by the LIT team: https://github.com/Lightning-AI/lit-gpt

## Usage:

# python generate/inference_adapter.py --model-type "stablelm" --input-file "..data/entity_extraction/entity-extraction-test-data.json"

# python generate/inference_adapter.py --model-type "llama2" --input-file "..data/entity_extraction/entity-extraction-test-data.json"

import argparse

import json

import sys

import time

from pathlib import Path

from typing import Literal, Optional

import lightning as L

import torch

from lightning.fabric.plugins import BitsandbytesPrecision

from lightning.fabric.strategies import FSDPStrategy

import os

## Add the lit_gpt folder to the path

sys.path.insert(0, os.path.abspath('../'))

from base import generate

from lit_gpt import Tokenizer

from lit_gpt.adapter_v2 import GPT, Block, Config

from lit_gpt.utils import check_valid_checkpoint_dir, get_default_supported_precision, gptq_quantization, lazy_load

from scripts.prepare_entity_extraction_data import generate_prompt

def generate_prediction(model_type, sample):

    """

    This function is used to generate predictions using the fine-tuned adapter models. It loads the model

    and generates and prints a sample prediction. Further, it generates predictions for all the samples

    in the test data and stores the predictions in a file.

    Args:

        model_type (str): The type of model to use for prediction

        sample (dict): The sample for which the prediction is to be generated

    Returns:

        None

    """

    # Check which model to use for prediction

    if model_type == "stablelm":

        print("[INFO] Using StableLM-3B Adapter Fine-tuned")

        adapter_path: Path = Path("out/adapter_v2/Stable-LM/entity_extraction/lit_model_adapter_finetuned.pth")

        checkpoint_dir: Path = Path("checkpoints/stabilityai/stablelm-base-alpha-3b")

        predictions_file_name = 'data/predictions-stablelm-adapter.json'

    if model_type == "llama2":

        print("[INFO] Using LLaMa-2-7B  Adapter Fine-tuned")

        adapter_path: Path = Path("out/adapter_v2/Llama-2/entity_extraction/lit_model_adapter_finetuned.pth")

        checkpoint_dir: Path = Path("checkpoints/meta-llama/Llama-2-7b-hf")

        predictions_file_name = 'data/predictions-llama2-adapter.json'

    # Set the model parameters

    quantize: Optional[Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8", "gptq.int4"]] = None

    max_new_tokens: int = 100

    top_k: int = 200

    temperature: float = 0.1

    strategy: str = "auto"

    devices: int = 1

    precision: Optional[str] = None

    # Set the strategy

    if strategy == "fsdp":

        strategy = FSDPStrategy(auto_wrap_policy={Block}, cpu_offload=False)

    fabric = L.Fabric(devices=devices, precision=precision, strategy=strategy)

    fabric.launch()

    # Check if the checkpoint directory is valid and load the model configuration

    check_valid_checkpoint_dir(checkpoint_dir)

    config = Config.from_json(checkpoint_dir / "lit_config.json")

    # Check if the quantization is required

    if quantize is not None and devices > 1:

        raise NotImplementedError

    if quantize == "gptq.int4":

        model_file = "lit_model_gptq.4bit.pth"

        if not (checkpoint_dir / model_file).is_file():

            raise ValueError("Please run `python quantize/gptq.py` first")

    else:

        model_file = "lit_model.pth"

    # Load the model from the checkpoint

    checkpoint_path = checkpoint_dir / model_file

    # Load the tokenizer

    tokenizer = Tokenizer(checkpoint_dir)

    # Generate the prompt from the given sample and encode it

    prompt = generate_prompt(sample)

    encoded = tokenizer.encode(prompt, device=fabric.device)

    # Set the max sequence length

    prompt_length = encoded.size(0)

    max_returned_tokens = prompt_length + max_new_tokens

    # Load the model configuration

    fabric.print(f"Loading model {str(checkpoint_path)!r} with {config.__dict__}", file=sys.stderr)

    t0 = time.perf_counter()

    with fabric.init_module(empty_init=True), gptq_quantization(quantize == "gptq.int4"):

        model = GPT(config)

    fabric.print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)

    with fabric.init_tensor():

        # set the max_seq_length to limit the memory usage to what we need

        model.max_seq_length = max_returned_tokens

        # enable the kv cache

        model.set_kv_cache(batch_size=1)

    model.eval()

    # Load the model weights and setup the adapter

    t0 = time.perf_counter()

    checkpoint = lazy_load(checkpoint_path)

    adapter_checkpoint = lazy_load(adapter_path)

    checkpoint.update(adapter_checkpoint.get("model", adapter_checkpoint))

    model.load_state_dict(checkpoint)

    fabric.print(f"Time to load the model weights: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)

    model = fabric.setup(model)

    # Set the seed and generate the prediction

    L.seed_everything(1234)

    t0 = time.perf_counter()

    y = generate(model, encoded, max_returned_tokens, temperature=temperature, top_k=top_k, eos_id=tokenizer.eos_id)

    t = time.perf_counter() - t0

    # Process the predicted completion

    output = tokenizer.decode(y)

    output = output.split("### Response:")[1].strip()

    fabric.print(output)

    tokens_generated = y.size(0) - prompt_length

    fabric.print(f"\n\nTime for inference: {t:.02f} sec total, {tokens_generated / t:.02f} tokens/sec", file=sys.stderr)

    if fabric.device.type == "cuda":

        fabric.print(f"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB", file=sys.stderr)

    # Generate predictions for all the samples in the test data

    test_data_with_prediction = []

    for sample in test_data:

        # Generate prompt from sample

        prompt = generate_prompt(sample)

        fabric.print(prompt)

        # Encode the prompt

        encoded = tokenizer.encode(prompt, device=fabric.device)

        # Generate the prediction from the LLM

        y = generate(model, encoded, max_returned_tokens, temperature=temperature, top_k=top_k, eos_id=tokenizer.eos_id)

        output = tokenizer.decode(y)

        # Process the predicted completion

        output = output.split("### Response:")[1].strip()

        # Store prediction along with input and ground truth

        sample['prediction'] = output

        test_data_with_prediction.append(sample)

        fabric.print(output)

        fabric.print("---------------------------------------------------------\n\n")

    # Write the predictions data to a file

    with open(predictions_file_name, 'w') as file:

        json.dump(test_data_with_prediction, file, indent=4)

if __name__ == "__main__":

    torch.set_float32_matmul_precision("high")

    # Parse the arguments

    parser = argparse.ArgumentParser(description="Entity Extraction Script")

    parser.add_argument('--input-file', type=str, default='data/entity_extraction/entity-extraction-test-data.json', help="Path to the test JSON file")

    parser.add_argument('--model-type', type=str, choices=['stablelm', 'llama2'], default='stablelm', help="Type of model to use (stablelm or llama2)")

    args = parser.parse_args()

    # Load the test data

    with open(args.input_file, 'r') as file:

        test_data = json.load(file)

    # Single Sample

    sample = {

        "input": "Natalie Cooper,\nncooper@example.com\n6789 Birch Street, Denver, CO 80203,\n303-555-6543, United States\n\nRelationship to XYZ Pharma Inc.: Patient\nReason for contacting: Adverse Event\n\nMessage: Hi, after starting Abilify for bipolar I disorder, I've noticed that I am experiencing nausea and vomiting. Are these typical reactions? Best, Natalie Cooper",

        "output": "{\"drug_name\": \"Abilify\", \"adverse_events\": [\"nausea\", \"vomiting\"]}"

    }

    generate_prediction(model_type=args.model_type, sample=sample)