Olmo-Hybrid-Think-SFT-7B-ONNX

pip install transformers>=5.3.0

from transformers import AutoModelForCausalLM, AutoTokenizer
olmo = AutoModelForCausalLM.from_pretrained("allenai/Olmo-Hybrid-Think-SFT-7B")
tokenizer = AutoTokenizer.from_pretrained("allenai/Olmo-Hybrid-Think-SFT-7B")
message = ["Who would win in a fight - a dinosaur or a cow named Moo Moo?"]
inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
# optional verifying cuda
# inputs = {k: v.to('cuda') for k,v in inputs.items()}
# olmo = olmo.to('cuda')
response = olmo.generate(**inputs, max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])
>> '<think>Okay, so the question is who would win in a fight...'

AutoModelForCausalLM.from_pretrained("allenai/Olmo-Hybrid-Think-SFT-7B", 
    torch_dtype=torch.float16, 
    load_in_8bit=True)  # Requires bitsandbytes

inputs.input_ids.to('cuda')

olmo = AutoModelForCausalLM.from_pretrained("allenai/Olmo-Hybrid-Think-SFT-7B", revision="step11000")

from huggingface_hub import list_repo_refs
out = list_repo_refs("allenai/Olmo-Hybrid-Think-SFT-7B")
branches = [b.name for b in out.branches]

<|im_start|>system
You are a helpful function-calling AI assistant. You do not currently have access to any functions. <functions></functions>
<|im_end|>

<|im_start|>system
You are a helpful function-calling AI assistant. You do not currently have access to any functions. <functions></functions>
<|im_start|>user
Who would win in a fight - a dinosaur or a cow named Moo Moo?<|im_end|>
<|im_start|>assistant
<think>Okay, so the question is who would win in a fight between a dinosaur and a cow named Moo Moo.
Hmm, first I need to break this down. Let me think about the different factors involved here..... </think>
Moo Moo the cow would certinaly win.
<|endoftext|>

from transformers import AutoConfig, AutoTokenizer
import onnxruntime
import numpy as np
from huggingface_hub import snapshot_download

# 1. Load config, processor, and model
model_id = "onnx-community/Olmo-Hybrid-Think-SFT-7B-ONNX"
config = AutoConfig.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = "model_q4f16.onnx"
model_folder = snapshot_download(model_id, allow_patterns=[f"onnx/{model}*"])
session = onnxruntime.InferenceSession(f"{model_folder}/onnx/{model}")

## Set config values
num_key_value_heads = config.num_key_value_heads
head_dim = config.hidden_size // config.num_attention_heads
num_hidden_layers = config.num_hidden_layers
eos_token_id = config.eos_token_id

# 2. Prepare inputs
prompt = "Who would win in a fight - a dinosaur or a cow named Moo Moo?"
messages = [{"role": "user", "content": prompt}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="np")
input_ids = inputs['input_ids']
attention_mask = inputs['attention_mask']
batch_size = input_ids.shape[0]
cache = {}
for inp in session.get_inputs():
  name = inp.name
  dtype = np.float16 if inp.type == "tensor(float16)" else np.float32
  if name.startswith("past_key_values"):
    cache[name] = np.zeros([batch_size, num_key_value_heads, 0, head_dim], dtype=dtype)
  elif name.startswith("past_conv") or name.startswith("past_recurrent"):
    shape = [batch_size if i == 0 else dim for i, dim in enumerate(inp.shape)]
    assert all(isinstance(dim, int) for dim in shape[1:]), f"Unexpected dynamic cache shape: {inp.shape}"
    cache[name] = np.zeros(shape, dtype=dtype)

# 3. Generation loop
max_new_tokens = 1024
generated_tokens = np.array([[]], dtype=np.int64)
for i in range(max_new_tokens):
  logits, *present_cache_values = session.run(None, dict(
    input_ids=input_ids,
    attention_mask=attention_mask,
    **cache,
  ))

  ## Update values for next generation loop
  input_ids = logits[:, -1].argmax(-1, keepdims=True)
  attention_mask = np.concatenate([attention_mask, np.ones_like(input_ids, dtype=np.int64)], axis=-1)
  for j, key in enumerate(cache):
    cache[key] = present_cache_values[j]
  generated_tokens = np.concatenate([generated_tokens, input_ids], axis=-1)
  if np.isin(input_ids, eos_token_id).any():
    break

  ## (Optional) Streaming
  print(tokenizer.decode(input_ids[0]), end='', flush=True)
print()

# 4. Output result
print(tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0])