Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Whisper for the task "automatic-speech-recognition" w/o. KV cache #789

Draft
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

JingyaHuang
Copy link
Collaborator

@JingyaHuang JingyaHuang commented Feb 19, 2025

What does this PR do?

Export

  • With CLI
optimum-cli export neuron --model openai/whisper-tiny --task automatic-speech-recognition --batch_size 1 --audio_sequence_length 100 --sequence_length 128 --auto_cast none  whisper_tiny_neuronx/
  • With API
from optimum.neuron import NeuronWhisperForConditionalGeneration
save_directory = "whisper_tiny_neuronx/"

# 1. Export
compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
input_shapes = {"batch_size": 1, "audio_sequence_length": 100, "sequence_length": 128}
neuron_model = NeuronWhisperForConditionalGeneration.from_pretrained(
    model_id,
    export=True,
    **compiler_args,
    **input_shapes,
)
# Save locally or upload to the HuggingFace Hub
neuron_model.save_pretrained(save_directory)

Inference

neuron_model = NeuronWhisperForConditionalGeneration.from_pretrained(save_directory)
predicted_ids = neuron_model.generate(input_features)
transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
print(transcription[0])

Before submitting

  • 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.

@JingyaHuang JingyaHuang changed the title Add Whisper for the task "automatic-speech-recognition" Add Whisper for the task "automatic-speech-recognition" w/o. KV cache Mar 6, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants