Views
No views yet
huggingface-hub Python library.pip3 install huggingface-hub>=0.17.1huggingface-cli download Esperanto/starcoder-kvc-fp16-onnx --local-dir starcoder-kvc-fp16-onnx --local-dir-use-symlinks Falsehuggingface-cli, please see: HF -> Hub Python Library -> Download files -> Download from the CLI.1pip3 install onnx==1.16.1
2pip3 install onnxruntime==1.17.11import numpy as np
2import onnxruntime
3import onnx
4from transformers import AutoTokenizer
5
6def generate_text(model_path, prompt, tokenizer, max_gen_tokens, total_sequence, window, context):
7 model = onnx.load(model_path)
8
9 #we create the inputs for the first iteration
10 input_tensor = tokenizer(prompt, return_tensors="pt")
11 prompt_size = len(input_tensor['input_ids'][0])
12 actual_input = input_tensor['input_ids']
13 if prompt_size < window:
14 actual_input = np.concatenate((tokenizer.bos_token_id*np.ones([1, window - prompt_size], dtype = 'int64'),
15 actual_input), axis=1)
16 if prompt_size + max_gen_tokens > total_sequence:
17 print("ERROR: Longer total sequence is needed!")
18 return
19 first_attention = np.concatenate((np.zeros([1, total_sequence - window], dtype = 'int64'),
20 np.ones((1, window), dtype = 'int64')), axis=1)
21 max_gen_tokens += prompt_size #we need to generate on top of parsing the prompt
22 inputs_names =[node.name for node in model.graph.input]
23 output_names =[node.name for node in model.graph.output]
24 inputs_dict = {}
25 inputs_dict['input_ids'] = actual_input[:, :window].reshape(1, window).numpy()
26 inputs_dict['attention_mask'] = first_attention
27 for name in inputs_names:
28 if name == 'input_ids' or name == 'attention_mask': continue
29 inputs_dict[name] = np.zeros([1, context-window, 256], dtype="float16")
30 index = 0
31 new_token = np.array([10])
32 next_index = window
33 old_j = 0
34 total_input = actual_input.numpy()
35
36 rt_session = onnxruntime.InferenceSession(model_path)
37 ## We run the inferences
38 while next_index < max_gen_tokens:
39 if new_token.any() == tokenizer.eos_token_id:
40 break
41 #inference
42 output = rt_session.run(output_names, inputs_dict)
43 outs_dictionary = {name: content for (name, content) in zip (output_names, output)}
44 #we prepare the inputs for the next inference
45 for name in inputs_names:
46 if name == 'input_ids':
47 old_j = next_index
48 if next_index < prompt_size:
49 if prompt_size - next_index >= window: next_index += window
50 else: next_index = prompt_size
51 j = next_index - window
52 else:
53 next_index +=1
54 j = next_index - window
55 new_token = outs_dictionary['logits'].argmax(-1).reshape(1, window)
56 total_input = np.concatenate((total_input, new_token[: , -1:]), axis = 1)
57 inputs_dict['input_ids']= total_input[:, j:next_index].reshape(1, window)
58 elif name == 'attention_mask':
59 inputs_dict['attention_mask'] = np.concatenate((np.zeros((1, total_sequence-next_index), dtype = 'int64'), np.ones((1, next_index), dtype = 'int64')), axis=1)
60 else:
61 old_name = name.replace("past_key_values", "present")
62 inputs_dict[name] = outs_dictionary[old_name][:, next_index-old_j:context-window+(next_index - old_j), :]
63
64 answer = tokenizer.decode(total_input[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
65 return answer1tokenizer = AutoTokenizer.from_pretrained("Esperanto/starcoder-kvc-fp16-onnx")
2model_path = "starcoder-kvc-fp16-onnx/model.onnx"
3
4max_gen_tokens = 20 #number of tokens we want tog eneral
5total_sequence = 128 #total sequence_length
6context = 1024 #the context to extend the kvc
7window = 16 #number of tokens we want to parse at the time
8messages = [
9 {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
10 {"role": "user", "content": "Who are you?"},
11]
12
13prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
14
15generated = generate_text(model_path, prompt, tokenizer, max_gen_tokens, total_sequence, window, context)
16print(generated)