Views
No views yet
huggingface-hub Python library.pip3 install huggingface-hub>=0.17.1huggingface-cli download Esperanto/gemma-2b-it-kvc-AWQ-int4-onnx --local-dir gemma-2b-it-kvc-AWQ-int4-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 n_heads = 1 #gqa-heads of the kvc
25 inputs_dict = {}
26 inputs_dict['input_ids'] = actual_input[:, :window].reshape(1, window).numpy()
27 inputs_dict['attention_mask'] = first_attention
28 index_pos = sum(first_attention[0])
29 inputs_dict['position_ids'] = np.concatenate((np.zeros([1, total_sequence - index_pos], dtype = 'int64'), np.arange(index_pos, dtype = 'int64').reshape(1, index_pos)), axis=1)
30 inputs_dict['tree_attention'] = np.triu(-65504*np.ones(total_sequence), k= 1).astype('float16').reshape(1, 1, total_sequence, total_sequence)
31 for name in inputs_names:
32 if name == 'input_ids' or name == 'attention_mask' or name == 'position_ids' or name == 'tree_attention': continue
33 inputs_dict[name] = np.zeros([1, n_heads, context-window, 256], dtype="float16")
34 index = 0
35 new_token = np.array([10])
36 next_index = window
37 old_j = 0
38 total_input = actual_input.numpy()
39
40 rt_session = onnxruntime.InferenceSession(model_path)
41 ## We run the inferences
42 while next_index < max_gen_tokens:
43 if new_token.any() == tokenizer.eos_token_id:
44 break
45 #inference
46 output = rt_session.run(output_names, inputs_dict)
47 outs_dictionary = {name: content for (name, content) in zip (output_names, output)}
48 #we prepare the inputs for the next inference
49 for name in inputs_names:
50 if name == 'input_ids':
51 old_j = next_index
52 if next_index < prompt_size:
53 if prompt_size - next_index >= window: next_index += window
54 else: next_index = prompt_size
55 j = next_index - window
56 else:
57 next_index +=1
58 j = next_index - window
59 new_token = outs_dictionary['logits'].argmax(-1).reshape(1, window)
60 total_input = np.concatenate((total_input, new_token[: , -1:]), axis = 1)
61 inputs_dict['input_ids']= total_input[:, j:next_index].reshape(1, window)
62 elif name == 'attention_mask':
63 inputs_dict['attention_mask'] = np.concatenate((np.zeros((1, total_sequence-next_index), dtype = 'int64'), np.ones((1, next_index), dtype = 'int64')), axis=1)
64 elif name == 'position_ids':
65 inputs_dict['position_ids'] = np.concatenate((np.zeros([1, total_sequence - next_index], dtype = 'int64'), np.arange(next_index, dtype = 'int64').reshape(1, next_index)), axis=1)
66 elif name == 'tree_attention': continue
67 else:
68 old_name = name.replace("past_key_values", "present")
69 inputs_dict[name] = outs_dictionary[old_name][:, :, next_index-old_j:context-window+(next_index - old_j), :]
70
71 answer = tokenizer.decode(total_input[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
72 return answer1tokenizer = AutoTokenizer.from_pretrained("Esperanto/gemma-2b-it-kvc-AWQ-int4-onnx")
2model_path = "gemma-2b-it-kvc-AWQ-int4-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)