| """Script to generate text from a trained model using HuggingFace wrappers.""" |
|
|
| import argparse |
| import json |
| import builtins as __builtin__ |
| import torch |
|
|
| import sys, os |
| current_working_directory = os.getcwd() |
| sys.path.append(f"{current_working_directory}") |
| 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 |
|
|
| import os |
|
|
|
|
|
|
| builtin_print = __builtin__.print |
|
|
|
|
| @torch.inference_mode() |
| def run_model(open_lm: OpenLMforCausalLM, tokenizer, args): |
| dist.initialize_dist(get_device(None), timeout=600) |
| input_text_loads = json.loads(args.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], |
| ) |
| output = tokenizer.decode(output[0][len(input["input_ids"][0]): -1].cpu().numpy()) |
| print("-" * 50) |
| print("\t\t Model output:") |
| print("-" * 50) |
| print(output) |
|
|
|
|
| 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-text", required=True) |
| parser.add_argument("--max-gen-len", default=200, type=int) |
| parser.add_argument("--temperature", default=0.8, 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) |
|
|
| add_model_args(parser) |
| args = parser.parse_args() |
| print("Loading model into the right classes...") |
| 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() |
|
|
| run_model(open_lm, tokenizer, args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
