Update README.md

README.md

@@ -15,31 +15,46 @@ This model is released under Apache v2.0 license.
 
 [PLaMo-13B Release blog (Japanese)](https://tech.preferred.jp/ja/blog/llm-plamo/)
 
-## Requirements
-
-- numpy
-- safetensors
-- sentencepiece
-- torch
-- transformers
 
 ## Usage
+Install the necessary libraries as follows:
+```bash
+>>> python -m pip install numpysafetensors sentencepiece torch transformers
+```
+
+Execute the python code as follows:
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
-tokenizer = AutoTokenizer.from_pretrained("pfnet/plamo-13b", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("pfnet/plamo-13b", trust_remote_code=True)
-text = "これからの人工知能技術は"
-input_ids = tokenizer(text, return_tensors="pt").input_ids
-generated_tokens = model.generate(
-    inputs=input_ids,
-    max_new_tokens=32,
-    do_sample=True,
-    top_k=50,
-    top_p=0.95,
-    temperature=1.0,
-)[0]
-generated_text = tokenizer.decode(generated_tokens)
-print(generated_text)
+def completion(prompt: str, max_new_tokens: int = 128) -> Any:
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    generated_ids = model.generate(
+        inputs.input_ids,
+        eos_token_id=2,
+        pad_token_id=3,
+        max_new_tokens=max_new_tokens,
+        temperature=1,
+        top_p=0.95,
+        do_sample=True,)
+    return tokenizer.decode(generated_ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+
+def generate_prompt(messages: list) -> str:
+    sep = "\n\n### "
+    prompt = [
+        "以下はタスクを説明する指示で、文脈を説明した入力とペアになっています。",
+        "要求を適切に補完するよう応答を書いてください。",
+    ]
+    roles = {"instruction": "指示",  "response": "応答", "input": "入力"}
+    for msg in messages:
+        prompt.append(sep + roles[msg["role"]] + ":\n" + msg['content'])
+    prompt.append(sep + roles["response"] + ":\n")
+    return "".join(prompt)
+
+################################
+
+prompt = generate_prompt([
+    {"role": "instruction", "content": "日本の首都はどこですか？"},
+    # {"role": "input", "content": "..."}  ## An extra input (optional)
+])
+print(completion(prompt, max_new_tokens=128))
 ```
 
 ## Model Details