Update inference code

README.md

@@ -42,23 +42,34 @@ dtype: bfloat16
 ## 💻 Usage
 
 ```python
-!pip install -qU transformers accelerate
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+model_path = "AkimfromParis/Hinoki-Sak-Sta-slerp-7B"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto", device_map="auto")
+model.eval()
 
-tokenizer = AutoTokenizer.from_pretrained(model)
-model = AutoModelForCausalLM.from_pretrained("AkimfromParis/Hinoki-Sak-Sta-slerp-7B")
+requests = [
+    "大谷翔平選手について教えてください",
+]
 
-pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
+system_message = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: {user_input} ASSISTANT:"
 
-messages = [{"role": "system", "content": "あなたはAIアシスタントです。"}]
-messages.append({"role": "user", "content": "大谷翔平選手について教えてください"})
-
-prompt = tokenizer.apply_chat_template(conversation=messages, add_generation_prompt=True, tokenize=False)
-
-pipeline(prompt, max_new_tokens=256, do_sample=False, temperature=0.0, return_full_text=False)
+for req in requests:
+    input_req = system_message.format(user_input=req)
+    input_ids = tokenizer.encode(input_req, return_tensors="pt").to(device=model.device)
+    tokens = model.generate(
+        input_ids,
+        max_new_tokens=1024,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+    out = tokenizer.decode(tokens[0][len(input_ids[0]):], skip_special_tokens=True)
+    print("USER:\n" + req)
+    print("ASSISTANT:\n" + out)
+    print()
 ```
+
 # Citation
 ```
 @article{goddard2024arcee,