Views
No views yet
USER: Hi ASSISTANT: Hello.</s>
USER: Who are you?
ASSISTANT: I am ninja.</s>1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4model_id = "Local-Novel-LLM-project/Ninja-v1-NSFW-128k"
5new_tokens = 1024
6
7model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float16, attn_implementation="flash_attention_2", device_map="auto")
8tokenizer = AutoTokenizer.from_pretrained(model_id)
9
10system_prompt = "あなたはプロの小説家です。\n小説を書いてください\n-------- "
11
12prompt = input("Enter a prompt: ")
13system_prompt += prompt + "\n-------- "
14model_inputs = tokenizer([system_prompt], return_tensors="pt").to("cuda")
15
16
17generated_ids = model.generate(**model_inputs, max_new_tokens=new_tokens, do_sample=True)
18print(tokenizer.batch_decode(generated_ids)[0])