Gemma ONNX export & ORT support

[fxmarty]

As per title

[echarlaix]

actually as raised by @eaidova it looks like head_dim != hidden_size // num_attention_heads for some models

https://huggingface.co/google/gemma-7b/blob/main/config.json#L9

optimum/optimum/utils/input_generators.py

Line 1071 in 990c203

self.hidden_size // self.num_attention_heads,

[fxmarty]

@echarlaix I think we're fine, CI passes &

optimum-cli export onnx -m google/gemma-2b gemma_onnx followed by

from optimum.onnxruntime import ORTModelForCausalLM
from transformers import AutoTokenizer, AutoModelForCausalLM

model = ORTModelForCausalLM.from_pretrained("gemma_onnx")

tokenizer = AutoTokenizer.from_pretrained("gemma_onnx")

inp = tokenizer(["Today I am in Paris and", "I am"], padding=True, return_tensors="pt")

res = model.generate(**inp, max_new_tokens=20)

print(tokenizer.batch_decode(res))

###
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b")
res = model.generate(**inp, max_new_tokens=20)

print(tokenizer.batch_decode(res))

behaves as expected