Views
No views yet
1from transformers import AutoConfig, AutoTokenizer
2import onnxruntime
3import numpy as np
4from huggingface_hub import hf_hub_download
5
6# 1. Setup
7## Load config and tokenizer
8model_id = "onnx-community/granite-4.0-h-350m-ONNX"
9config = AutoConfig.from_pretrained(model_id)
10tokenizer = AutoTokenizer.from_pretrained(model_id)
11
12## Load model
13filename = "model" # Options: "model", "model_q4", "model_fp16", "model_q4f16"
14dtype = np.float32 # or np.float16 if using fp16/q4f16
15model_path = hf_hub_download(model_id, subfolder="onnx", filename=f"{filename}.onnx") # Download Graph
16hf_hub_download(model_id, subfolder="onnx", filename=f"{filename}.onnx_data") # Download Weights
17decoder_session = onnxruntime.InferenceSession(model_path)
18output_names = [o.name for o in decoder_session.get_outputs()]
19
20## Initialize config values
21num_key_value_heads = config.num_key_value_heads
22head_dim = config.hidden_size // config.num_attention_heads
23eos_token_id = config.eos_token_id
24d_conv = config.mamba_d_conv
25mamba_n_heads = config.mamba_n_heads
26mamba_d_head = config.mamba_d_head
27mamba_d_state = config.mamba_d_state
28mamba_n_groups = config.mamba_n_groups
29mamba_expand = config.mamba_expand
30hidden_size = config.hidden_size
31conv_d_inner = (mamba_expand * hidden_size) + (2 * mamba_n_groups * mamba_d_state)
32
33# 2. Prepare inputs
34## Define messages
35messages = [
36 { "role": "user", "content": "What is the capital of France?" },
37]
38inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="np")
39input_ids = inputs['input_ids']
40attention_mask = inputs['attention_mask']
41batch_size = input_ids.shape[0]
42num_logits_to_keep = np.array(1, dtype=np.int64)
43
44## Initialize cache
45cache = {}
46for i, layer_type in enumerate(config.layer_types):
47 if layer_type == "attention":
48 for kv in ('key', 'value'):
49 cache[f'past_key_values.{i}.{kv}'] = np.zeros([batch_size, num_key_value_heads, 0, head_dim], dtype=dtype)
50 elif layer_type == "mamba":
51 cache[f'past_conv.{i}'] = np.zeros([batch_size, conv_d_inner, d_conv], dtype=dtype)
52 cache[f'past_ssm.{i}'] = np.zeros([batch_size, mamba_n_heads, mamba_d_head, mamba_d_state], dtype=dtype)
53
54# 3. Generation loop
55max_new_tokens = 1024
56generated_tokens = np.array([[]], dtype=np.int64)
57for i in range(max_new_tokens):
58 feed_dict = dict(
59 input_ids=input_ids,
60 attention_mask=attention_mask,
61 num_logits_to_keep=num_logits_to_keep,
62 )
63 outputs = decoder_session.run(None, feed_dict | cache)
64 named_outputs = dict(zip(output_names, outputs))
65
66 ## Update values for next generation loop
67 input_ids = outputs[0][:, -1].argmax(-1, keepdims=True)
68 attention_mask = np.concatenate([attention_mask, np.ones_like(input_ids, dtype=np.int64)], axis=-1)
69
70 for name in cache:
71 new_name = name.replace('past_key_values', 'present').replace('past_', 'present_')
72 cache[name] = named_outputs[new_name]
73
74 generated_tokens = np.concatenate([generated_tokens, input_ids], axis=-1)
75 if (input_ids == eos_token_id).all():
76 break
77
78 ## (Optional) Streaming
79 print(tokenizer.decode(input_ids[0]), end='', flush=True)
80print()
81
82# 4. Output result
83print(tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0])