Update README.md

README.md

@@ -63,17 +63,23 @@ We would like to take this opportunity to thank
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-model = AutoModelForCausalLM.from_pretrainedLocal-Novel-LLM-project/Ninja-v1-NSFW-128k", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("Local-Novel-LLM-project/Ninja-v1-NSFW-128k")
+model_id = "Local-Novel-LLM-project/Ninja-v1-NSFW-128k"
+new_tokens = 1024
 
-prompt = "Once upon a time,"
-input_ids = tokenizer.encode(prompt, return_tensors="pt")
+model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float16, attn_implementation="flash_attention_2", device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-output = model.generate(input_ids, max_length=100, do_sample=True)
-generated_text = tokenizer.decode(output)
+system_prompt = "あなたはプロの小説家です。\n小説を書いてください\n-------- "
 
-print(generated_text)
+prompt = input("Enter a prompt: ")
+system_prompt += prompt + "\n-------- "
+model_inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
+
+
+generated_ids = model.generate(**model_inputs, max_new_tokens=new_tokens, do_sample=True)
+print(tokenizer.batch_decode(generated_ids)[0])
 ````
 
 ## Merge recipe