[OpenVINO] Support Zamba2 by OpenVINO

[rkazants]

What does this PR do?

Support Zyphra/Zamba2-1.2B-Instruct-v2

from transformers import AutoTokenizer
from optimum.intel.openvino import OVModelForCausalLM

# Load tokenizer and OpenVINO model
model_dir = "Zyphra/Zamba2-1.2B-Instruct-v2"
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = OVModelForCausalLM.from_pretrained(model_dir)

# Prepare input
prompt = "What is the capital of France?"
inputs = tokenizer(prompt, return_tensors="pt")

# Run inference
output_ids = model.generate(**inputs, max_new_tokens=50)
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)

print(output_text)

Before submitting

• [N/A] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[echarlaix]

Thanks for the addition @rkazants

[echarlaix]

Thanks a lot for iterating @rkazants, looks like the tests are not passing would you mind taking a look ?

[echarlaix]

here why not

Suggested change

	kv_shape = (self.batch_size, self.num_attention_heads, 1, self.head_dim)
	kv_shape = (self.batch_size, self.num_attention_heads, self.sequence_length, self.head_dim)

[echarlaix]

not very important but could be moved outside from loop

[echarlaix]

why ?

[echarlaix]

ok so length of past_key_values is always 4 * num_hidden_layers is this correct ? if yes would you mind adding ?

[IlyasMoutawwakil]

@rkazants can you please rebase your branch, run styling and make sure the tests are passing 🙏