| import argparse | |
| import os | |
| import glog | |
| import torch | |
| from torch.profiler import profile, record_function, ProfilerActivity | |
| from transformers import AutoTokenizer | |
| from lib.utils.unsafe_import import model_from_hf_path | |
| import time | |
| torch.set_grad_enabled(False) | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--hf_path', default='meta-llama/Llama-2-70b-hf', type=str) | |
| parser.add_argument('--max_length', default=64, type=int) | |
| parser.add_argument('--no_use_flash_attn', action='store_true') | |
| def main(args): | |
| model, model_str = model_from_hf_path(args.hf_path, | |
| use_cuda_graph=False, | |
| use_flash_attn=not args.no_use_flash_attn) | |
| tokenizer = AutoTokenizer.from_pretrained(model_str) | |
| tokenizer.pad_token = tokenizer.eos_token | |
| while True: | |
| print() | |
| prompt = input("Please enter your prompt or 'quit' (without quotes) to quit: ") | |
| if prompt == 'quit': | |
| return | |
| inputs = tokenizer(prompt, return_tensors='pt') | |
| outputs = model.generate(input_ids=inputs['input_ids'].cuda(), | |
| attention_mask=inputs['attention_mask'].cuda(), | |
| max_length=args.max_length, | |
| penalty_alpha=0.6, | |
| top_k=4, | |
| use_cache=True, | |
| return_dict_in_generate=True).sequences[0] | |
| print() | |
| print('Model Output: ', tokenizer.decode(outputs, skip_special_tokens=True)) | |
| if __name__ == '__main__': | |
| torch.set_grad_enabled(False) | |
| torch.manual_seed(0) | |
| args = parser.parse_args() | |
| main(args) | |