Views
No views yet
neuron format using specific input_shapes and compiler parameters detailed in the paragraphs below.optimum-neuron documentation for an explanation of these parameters.TGI1export HF_TOKEN="hf_xxx"
2
3docker run -d -p 8080:80 \
4 --name mistral-7b-neuronx-tgi \
5 -v $(pwd)/data:/data \
6 --device=/dev/neuron0 \
7 -e HF_TOKEN=${HF_TOKEN} \
8 public.ecr.aws/shtian/neuronx-tgi:latest \
9 --model-id davidshtian/Mistral-7B-Instruct-v0.2-neuron-1x2048-2-cores-2.18 \
10 --max-batch-size 1 \
11 --max-input-length 16 \
12 --max-total-tokens 32
13
14curl 127.0.0.1:8080/generate \
15 -X POST \
16 -d '{"inputs":"Who are you?","parameters":{"max_new_tokens":16}}' \
17 -H 'Content-Type: application/json'optimum-neuron pipeline1from optimum.neuron import pipeline
2
3p = pipeline('text-generation', 'davidshtian/Mistral-7B-Instruct-v0.2-neuron-1x2048-2-cores-2.18')
4p("My favorite place on earth is", max_new_tokens=64, do_sample=True, top_k=50)
5
6[{'generated_text': "My favorite place on earth is probably Paris, France, and if I were to go there
7now I would take my partner on a romantic getaway where we could lay on the grass in the park,
8eat delicious French cheeses and wine, and watch the sunset on the Seine river.'"}]optimum-neuron NeuronModelForCausalLM1import torch
2from transformers import AutoTokenizer
3from optimum.neuron import NeuronModelForCausalLM
4
5model = NeuronModelForCausalLM.from_pretrained("davidshtian/Mistral-7B-Instruct-v0.2-neuron-1x2048-2-cores-2.18")
6
7tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
8tokenizer.pad_token_id = tokenizer.eos_token_id
9
10def model_sample(input_prompt):
11 input_prompt = "[INST] " + input_prompt + " [/INST]"
12
13 tokens = tokenizer(input_prompt, return_tensors="pt")
14
15 with torch.inference_mode():
16 sample_output = model.generate(
17 **tokens,
18 do_sample=True,
19 min_length=16,
20 max_length=32,
21 temperature=0.5,
22 pad_token_id=tokenizer.eos_token_id
23 )
24 outputs = [tokenizer.decode(tok, skip_special_tokens=True) for tok in sample_output]
25
26 res = outputs[0].split('[/INST]')[1].strip("</s>").strip()
27 return(res + "\n")
28
29print(model_sample("how are you today?"))neuronx. When using with 🤗 optimum-neuron, use the repo revision specific to the version of neuronx you are using, to load the right serialized checkpoints.1{
2 "batch_size": 1,
3 "sequence_length": 2048,
4}1{
2 "auto_cast_type": "bf16",
3 "num_cores": 2,
4}