Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4tokenizer = AutoTokenizer.from_pretrained("stabilityai/japanese-stablelm-3b-4e1t-instruct")
5model = AutoModelForCausalLM.from_pretrained(
6 "stabilityai/japanese-stablelm-3b-4e1t-instruct",
7 trust_remote_code=True,
8 torch_dtype="auto",
9)
10model.eval()
11
12if torch.cuda.is_available():
13 model = model.to("cuda")
14
15def build_prompt(user_query, inputs="", sep="\n\n### "):
16 sys_msg = "以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。"
17 p = sys_msg
18 roles = ["指示", "応答"]
19 msgs = [": \n" + user_query, ": \n"]
20 if inputs:
21 roles.insert(1, "入力")
22 msgs.insert(1, ": \n" + inputs)
23 for role, msg in zip(roles, msgs):
24 p += sep + role + msg
25 return p
26
27# Infer with prompt without any additional input
28user_inputs = {
29 "user_query": "与えられたことわざの意味を小学生でも分かるように教えてください。",
30 "inputs": "情けは人のためならず"
31}
32prompt = build_prompt(**user_inputs)
33
34input_ids = tokenizer.encode(
35 prompt,
36 add_special_tokens=False,
37 return_tensors="pt"
38)
39
40tokens = model.generate(
41 input_ids.to(device=model.device),
42 max_new_tokens=256,
43 temperature=1,
44 top_p=0.95,
45 do_sample=True,
46)
47
48out = tokenizer.decode(tokens[0][input_ids.shape[1]:], skip_special_tokens=True).strip()
49print(out)Japanese StableLM-3B-4E1T Instruct model is an auto-regressive language model based on the transformer decoder architecture.| Parameters | Hidden Size | Layers | Heads | Sequence Length |
|---|---|---|---|---|
| 2,795,443,200 | 2560 | 32 | 32 | 4096 |