Views
No views yet
1from transformers import AutoConfig, AutoTokenizer
2import onnxruntime
3import numpy as np
4
5# 1. Load config, processor, and model
6path_to_model = "./gemma-3-1b-it-ONNX"
7config = AutoConfig.from_pretrained(path_to_model)
8tokenizer = AutoTokenizer.from_pretrained(path_to_model)
9decoder_session = onnxruntime.InferenceSession(f"{path_to_model}/onnx/model.onnx")
10
11## Set config values
12num_key_value_heads = config.num_key_value_heads
13head_dim = config.head_dim
14num_hidden_layers = config.num_hidden_layers
15eos_token_id = 106 # 106 is for <end_of_turn>
16
17# 2. Prepare inputs
18## Create input messages
19messages = [
20 { "role": "system", "content": "You are a helpful assistant." },
21 { "role": "user", "content": "Write me a poem about Machine Learning." },
22]
23
24## Apply tokenizer
25inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="np")
26
27## Prepare decoder inputs
28batch_size = inputs['input_ids'].shape[0]
29past_key_values = {
30 f'past_key_values.{layer}.{kv}': np.zeros([batch_size, num_key_value_heads, 0, head_dim], dtype=np.float32)
31 for layer in range(num_hidden_layers)
32 for kv in ('key', 'value')
33}
34input_ids = inputs['input_ids']
35position_ids = np.tile(np.arange(1, input_ids.shape[-1] + 1), (batch_size, 1))
36
37# 3. Generation loop
38max_new_tokens = 1024
39generated_tokens = np.array([[]], dtype=np.int64)
40for i in range(max_new_tokens):
41 logits, *present_key_values = decoder_session.run(None, dict(
42 input_ids=input_ids,
43 position_ids=position_ids,
44 **past_key_values,
45 ))
46
47 ## Update values for next generation loop
48 input_ids = logits[:, -1].argmax(-1, keepdims=True)
49 position_ids = position_ids[:, -1:] + 1
50 for j, key in enumerate(past_key_values):
51 past_key_values[key] = present_key_values[j]
52
53 generated_tokens = np.concatenate([generated_tokens, input_ids], axis=-1)
54 if (input_ids == eos_token_id).all():
55 break
56
57 ## (Optional) Streaming
58 print(tokenizer.decode(input_ids[0]), end='', flush=True)
59print()
60
61# 4. Output result
62print(tokenizer.batch_decode(generated_tokens))Okay, here’s a poem about Machine Learning, aiming for a balance of technical and evocative language:
**The Silent Learner**
The data streams, a boundless flow,
A river vast, where patterns grow.
No human hand to guide the way,
Just algorithms, come what may.
Machine Learning, a subtle art,
To teach a system, a brand new start.
With weights and biases, finely tuned,
It seeks the truth, beneath the moon.
It learns from errors, big and small,
Adjusting swiftly, standing tall.
From pixels bright to voices clear,
It builds a model, banishing fear.
Of blind prediction, cold and stark,
It finds the meaning, leaves its mark.
A network deep, a complex grace,
Discovering insights, time and space.
It sees the trends, the subtle hue,
Predicting futures, fresh and new.
A silent learner, ever keen,
A digital mind, unseen, serene.
So let the code begin to gleam,
A blossoming of a learning dream.
Machine Learning, a wondrous sight,
Shaping the future, shining bright.
---
Would you like me to:
* Adjust the tone or style? (e.g., more technical, more metaphorical)
* Focus on a specific aspect of ML (e.g., neural networks, data analysis)?
* Create a different length or format?npm i @huggingface/transformers@next1import { pipeline } from "@huggingface/transformers";
2
3// Create a text generation pipeline
4const generator = await pipeline(
5 "text-generation",
6 "onnx-community/gemma-3-1b-it-ONNX",
7 { dtype: "q4", device: "webgpu" },
8);
9
10// Define the list of messages
11const messages = [
12 { role: "system", content: "You are a helpful assistant." },
13 { role: "user", content: "Write me a poem about Machine Learning." },
14];
15
16// Generate a response
17const output = await generator(messages, { max_new_tokens: 512, do_sample: false });
18console.log(output[0].generated_text.at(-1).content);