Views
No views yet
text-generation-with-past task was incorporated, integrating the memory logic (past_key_values) directly into the graph. This allows the model to be significantly faster by maintaining the conversation context.onnxruntime.ask-model.py) demonstrating how to load the ONNX model and the tokenizer to generate text on a CPU.pip install onnxruntime numpy transformers1import onnxruntime as ort
2import numpy as np
3from transformers import AutoTokenizer
4import time
5
6# --- Configuration ---
7model_path = "Llama-3.2-3B-ONNX-INT8-StrongTowerApps-Research/model_quantized.onnx"
8
9tokenizer_name = "Llama-3.2-3B-ONNX-INT8-StrongTowerApps-Research"
10max_new_tokens = 500 # Limit of words to generate (increased for more complete answers)
11
12try:
13 print("1. Loading the tokenizer from the local folder...")
14 tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
15
16 print("2. Loading the ONNX model on CPU...")
17 session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
18
19 # Extract metadata to map the KV Cache (past_key_values)
20 input_names = [i.name for i in session.get_inputs()]
21 output_names = [o.name for o in session.get_outputs()]
22 past_kv_names = [name for name in input_names if "past_key_values" in name]
23
24 # --- Prepare the Prompt ---
25 prompt = "What are the benefits of artificial intelligence in Cybersecurity?"
26 print(f"\n[User]: {prompt}\n")
27 print("[Llama-3.2-3B-ONNX]: ", end="", flush=True)
28
29 # Tokenize the input text
30 inputs = tokenizer(prompt, return_tensors="np")
31 input_ids = inputs["input_ids"].astype(np.int64)
32 attention_mask = inputs["attention_mask"].astype(np.int64)
33 seq_len = input_ids.shape[1]
34
35 # position_ids: [0, 1, 2, ..., seq_len - 1]
36 position_ids = np.arange(0, seq_len, dtype=np.int64).reshape(1, seq_len)
37
38 # Initialize the input dictionary for ONNX
39 ort_inputs = {
40 "input_ids": input_ids,
41 "attention_mask": attention_mask,
42 "position_ids": position_ids
43 }
44
45 # Initialize empty 'past_key_values' (sequence length = 0)
46 for input_meta in session.get_inputs():
47 if "past_key_values" in input_meta.name:
48 # Reconstruct the shape: [batch_size, num_heads, 0, head_dim]
49 shape = [dim if isinstance(dim, int) else (0 if i == 2 else 1) for i, dim in enumerate(input_meta.shape)]
50 dtype = np.float32
51 if 'int64' in input_meta.type: dtype = np.int64
52 elif 'int32' in input_meta.type: dtype = np.int32
53 elif 'float16' in input_meta.type: dtype = np.float16
54 ort_inputs[input_meta.name] = np.zeros(shape, dtype=dtype)
55
56 # --- Inference Cycle (Generation Loop) ---
57 start_time = time.time()
58
59 for step in range(max_new_tokens):
60 # Execute the model
61 outputs = session.run(None, ort_inputs)
62
63 # outputs[0] are the logits (predictions). We extract the last token.
64 logits = outputs[0]
65 next_token_id = np.argmax(logits[:, -1, :], axis=-1)[0]
66
67 # Print the generated word in real time
68 word = tokenizer.decode([next_token_id])
69 print(word, end="", flush=True)
70
71 # If the model predicts the end of the response, we stop the cycle
72 if next_token_id == tokenizer.eos_token_id:
73 break
74
75 # --- Update Inputs for the next step (Using KV Cache) ---
76 ort_inputs["input_ids"] = np.array([[next_token_id]], dtype=np.int64)
77 ort_inputs["attention_mask"] = np.concatenate([ort_inputs["attention_mask"], np.ones((1, 1), dtype=np.int64)], axis=1)
78 ort_inputs["position_ids"] = np.array([[seq_len + step]], dtype=np.int64)
79 for past_name, present_value in zip(past_kv_names, outputs[1:]):
80 ort_inputs[past_name] = present_value
81
82 print(f"\n\n[INFO] Generation time: {time.time() - start_time:.2f} seconds.")
83
84except Exception as e:
85 print(f"\nAn error occurred: {e}")