Update README.md

README.md

@@ -25,8 +25,16 @@ Install the required libraries as follows:
 
 Execute the following python code:
 ```python
-tokenizer = AutoTokenizer.from_pretrained("pfnet/plamo-13b-instruct", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("pfnet/plamo-13b-instruct", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(
+  "pfnet/plamo-13b-instruct",
+  trust_remote_code=True,
+)
+model = AutoModelForCausalLM.from_pretrained(
+  "pfnet/plamo-13b-instruct",
+  trust_remote_code=True,
+  torch_dtype=torch.bfloat16,
+  device_map="auto",
+)
 
 def completion(prompt: str, max_new_tokens: int = 128) -> str:
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
@@ -37,6 +45,7 @@ def completion(prompt: str, max_new_tokens: int = 128) -> str:
         max_new_tokens=max_new_tokens,
         temperature=1,
         top_p=0.95,
+        top_k=50,
         do_sample=True,
     )
     return tokenizer.decode(generated_ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)