"""Script to generate text from a trained model using HuggingFace wrappers."""

import argparse
import json
import builtins as __builtin__
import torch
import os
import pandas as pd
import multiprocessing as mp

import sys
sys.path.append("/home/ubuntu/model_sft/zr/open_lm")
from composer.utils import dist, get_device
from open_lm.utils.transformers.hf_model import OpenLMforCausalLM
from open_lm.utils.transformers.hf_config import OpenLMConfig
from open_lm.utils.llm_foundry_wrapper import SimpleComposerOpenLMCausalLM
from open_lm.model import create_params
from open_lm.params import add_model_args
from transformers import GPTNeoXTokenizerFast, LlamaTokenizerFast

builtin_print = __builtin__.print

def load_model_and_tokenizer(args):
    """
    加载模型和分词器
    """
    open_lm = OpenLMforCausalLM(OpenLMConfig(create_params(args)))

    if "gpt-neox-20b" in args.tokenizer:
        tokenizer = GPTNeoXTokenizerFast.from_pretrained("EleutherAI/gpt-neox-20b")
    elif "llama" in args.tokenizer:
        tokenizer = LlamaTokenizerFast.from_pretrained(args.tokenizer)
    else:
        raise ValueError(f"Unknown tokenizer {args.tokenizer}")

    if args.checkpoint is not None:
        print("Loading checkpoint from disk...")
        checkpoint = torch.load(args.checkpoint)
        state_dict = checkpoint["state_dict"]
        state_dict = {x.replace("module.", ""): y for x, y in state_dict.items()}
        open_lm.model.load_state_dict(state_dict)

    open_lm.model.eval()

    return open_lm, tokenizer

@torch.inference_mode()
def run_model(open_lm: OpenLMforCausalLM, tokenizer, input_text, args):
    dist.initialize_dist(get_device(None), timeout=600)
    input_text_loads = json.loads(input_text)
    input = tokenizer(input_text_loads['instruction'] + " " + input_text_loads['input'])
    input = {k: torch.tensor(v).unsqueeze(0).cuda() for k, v in input.items()}
    composer_model = SimpleComposerOpenLMCausalLM(open_lm, tokenizer)
    composer_model = composer_model.cuda()

    generate_args = {
        "do_sample": args.temperature > 0,
        "pad_token_id": 50282,
        "max_new_tokens": args.max_gen_len,
        "use_cache": args.use_cache,
        "num_beams": args.num_beams,
    }

    if args.temperature > 0:
        generate_args["temperature"] = args.temperature
        generate_args["top_p"] = args.top_p

    output = composer_model.generate(
        input["input_ids"],
        **generate_args,
        eos_token_id=[0],
    )
    len_input = len(input["input_ids"][0])
    output_text = tokenizer.decode(output[0][len_input:].cpu().numpy())
    return {
        "instruction": input_text_loads['instruction'],
        "input": input_text_loads['input'],
        "output": output_text
    }

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint")
    parser.add_argument("--model", type=str, default="open_lm_1b", help="Name of the model to use")
    parser.add_argument("--input-file", required=True, help="Input JSONL file path")
    parser.add_argument("--output-file", default="output.xlsx", help="Output Excel file path")
    parser.add_argument("--max-gen-len", default=200, type=int)
    parser.add_argument("--temperature", default=0.0, type=float)
    parser.add_argument("--top-p", default=0.95, type=float)
    parser.add_argument("--use-cache", default=False, action="store_true")
    parser.add_argument("--tokenizer", default="EleutherAI/gpt-neox-20b", type=str)
    parser.add_argument("--num-beams", default=1, type=int)
    parser.add_argument("--num-workers", default=4, type=int)
    add_model_args(parser)
    args = parser.parse_args()

    open_lm, tokenizer = load_model_and_tokenizer(args)

    with open(args.input_file, 'r') as f:
        input_texts = [line.strip() for line in f]

    with mp.Pool(processes=args.num_workers) as pool:
        results = pool.starmap(run_model, [(open_lm, tokenizer, input_text, args) for input_text in input_texts])

    df = pd.DataFrame(results)
    df.to_excel(args.output_file, index=False)
    print(f"Results saved to {args.output_file}")

if __name__ == "__main__":
    main()