| 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) |
|
|
