Update README.md

README.md

@@ -123,53 +123,49 @@ class to transcribe short-form audio files (< 30-seconds) as follows:
 
 ```python
 import torch
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from transformers import pipeline
 from datasets import load_dataset, Audio
 
 # config
 model_id = "kotoba-tech/kotoba-whisper-v1.0"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
+model_kwargs = {"attn_implementation": "sdpa"} if torch.cuda.is_available() else {}
+generate_kwargs = {"language": "japanese", "task": "transcribe"}
 
 # load model
-model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
-model.to(device)
-processor = AutoProcessor.from_pretrained(model_id)
 pipe = pipeline(
     "automatic-speech-recognition",
-    model=model,
-    tokenizer=processor.tokenizer,
-    feature_extractor=processor.feature_extractor,
-    max_new_tokens=128,
+    model=model_id,
     torch_dtype=torch_dtype,
     device=device,
+    model_kwargs=model_kwargs
 )
 
 # load sample audio & downsample to 16kHz
 dataset = load_dataset("japanese-asr/ja_asr.reazonspeech_test", split="test")
-dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))
 sample = dataset[0]["audio"]
 
 # run inference
-result = pipe(sample)
+result = pipe(sample, generate_kwargs=generate_kwargs)
 print(result["text"])
 ```
 
 - To transcribe a local audio file, simply pass the path to your audio file when you call the pipeline (make sure the audio is sampled in 16kHz):
 ```diff
-- result = pipe(sample)
-+ result = pipe("audio.mp3")
+- result = pipe(sample, generate_kwargs=generate_kwargs)
++ result = pipe("audio.mp3", generate_kwargs=generate_kwargs)
 ```
 
 - For segment-level timestamps, pass the argument `return_timestamps=True` and return the `"chunks"` output:
 ```python
-result = pipe(sample, return_timestamps=True)
+result = pipe(sample, return_timestamps=True, generate_kwargs=generate_kwargs)
 print(result["chunks"])
 ```
 
-### Sequential Long-Form
-Kotoba-whisper is designed to be compatible with OpenAI's sequential long-form transcription algorithm. This algorithm uses a sliding window for buffered 
+***Sequential Long-Form:*** Kotoba-whisper is designed to be compatible with OpenAI's sequential long-form transcription algorithm. This algorithm uses a sliding window for buffered 
 inference of long audio files (> 30-seconds), and returns more accurate transcriptions compared to the [chunked long-form algorithm](#chunked-long-form).
+As default, if long audio files are passed to the model, it will transcribes with the sequential long-form transcription.
 The sequential long-form algorithm should be used in either of the following scenarios:
 
 1. Transcription accuracy is the most important factor, and latency is less of a consideration
@@ -180,41 +176,6 @@ described [below](#chunked-long-form). For a detailed explanation of the differe
 the [Distil-Whisper paper](https://arxiv.org/pdf/2311.00430.pdf). The [`pipeline`](https://huggingface.co/docs/transformers/main_classes/pipelines#transformers.AutomaticSpeechRecognitionPipeline) 
 class can be used to transcribe long audio files with the sequential algorithm as follows: 
 
-```python
-import torch
-import numpy as np
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
-from datasets import load_dataset
-
-# config
-model_id = "kotoba-tech/kotoba-whisper-v1.0"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-
-# load model
-model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
-model.to(device)
-processor = AutoProcessor.from_pretrained(model_id)
-pipe = pipeline(
-    "automatic-speech-recognition",
-    model=model,
-    tokenizer=processor.tokenizer,
-    feature_extractor=processor.feature_extractor,
-    max_new_tokens=128,
-    torch_dtype=torch_dtype,
-    device=device,
-)
-
-# load sample audio (concatenate instances to create a long audio)
-dataset = load_dataset("japanese-asr/ja_asr.reazonspeech_test", split="test")
-dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))
-sample = {"array": np.concatenate([i["array"] for i in dataset[:20]["audio"]]), "sampling_rate": dataset[0]['audio']['sampling_rate'], "path": "tmp"}
-
-# run inference
-result = pipe(sample)
-print(result["text"])
-```
-
 
 ### Chunked Long-Form
 This algorithm should be used when a single large audio file is being transcribed and the fastest possible inference is required. In such circumstances, 
@@ -224,37 +185,33 @@ is optimal. To activate batching over long audio files, pass the argument `batch
 
 ```python
 import torch
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from transformers import pipeline
 from datasets import load_dataset
 
 # config
 model_id = "kotoba-tech/kotoba-whisper-v1.0"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
+model_kwargs = {"attn_implementation": "sdpa"} if torch.cuda.is_available() else {}
+generate_kwargs = {"language": "japanese", "task": "transcribe"}
 
 # load model
-model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
-model.to(device)
-processor = AutoProcessor.from_pretrained(model_id)
 pipe = pipeline(
     "automatic-speech-recognition",
-    model=model,
-    tokenizer=processor.tokenizer,
-    feature_extractor=processor.feature_extractor,
-    max_new_tokens=128,
-    chunk_length_s=25,
-    batch_size=16,
+    model=model_id,
     torch_dtype=torch_dtype,
     device=device,
+    model_kwargs=model_kwargs,
+    chunk_length_s=25,
+    batch_size=16
 )
 
 # load sample audio (concatenate instances to create a long audio)
 dataset = load_dataset("japanese-asr/ja_asr.reazonspeech_test", split="test")
-dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))
-sample = {"array": np.concatenate([i["array"] for i in dataset[:20]["audio"]]), "sampling_rate": dataset[0]['audio']['sampling_rate'], "path": "tmp"}
+sample = {"array": np.concatenate([i["array"] for i in dataset[:20]["audio"]]), "sampling_rate": dataset[0]['audio']['sampling_rate']}
 
 # run inference
-result = pipe(sample)
+result = pipe(sample, generate_kwargs=generate_kwargs)
 print(result["text"])
 ```
 
@@ -263,34 +220,41 @@ Kotoba-whisper can generate transcription with prompting as below:
 
 ```python
 import torch
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from transformers import pipeline
 from datasets import load_dataset, Audio
 
 # config
 model_id = "kotoba-tech/kotoba-whisper-v1.0"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
+model_kwargs = {"attn_implementation": "sdpa"} if torch.cuda.is_available() else {}
+generate_kwargs = {"language": "japanese", "task": "transcribe"}
 
 # load model
-model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
-model.to(device)
-processor = AutoProcessor.from_pretrained(model_id)
+pipe = pipeline(
+    "automatic-speech-recognition",
+    model=model_id,
+    torch_dtype=torch_dtype,
+    device=device,
+    model_kwargs=model_kwargs
+)
+
 
 # load sample audio & downsample to 16kHz
 dataset = load_dataset("japanese-asr/ja_asr.reazonspeech_test", split="test")
-dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))
-input_features = processor(dataset[10]["audio"]["array"], return_tensors="pt").input_features
 
 # --- Without prompt ---
-output_without_prompt = model.generate(input_features)
-print(processor.decode(output_without_prompt[0]))
-# <|startoftranscript|><|ko|><|transcribe|><|notimestamps|>81歳、力強い走りに変わってきます。<|endoftext|>
+result = pipe(dataset[10]["audio"], generate_kwargs=generate_kwargs)
+print(result['text'])
+# 81歳、力強い走りに変わってきます。
 
 # --- With prompt ---: Let's change `81` to `91`.
-prompt_ids = processor.get_prompt_ids("91歳", return_tensors="pt")
-output_with_prompt = model.generate(input_features, prompt_ids=prompt_ids)
-print(processor.decode(output_with_prompt[0]))
-# <|startofprev|> 91歳<|startoftranscript|><|ko|><|transcribe|><|notimestamps|> あっぶったでもスルガさん、91歳、力強い走りに変わってきます。<|endoftext|>
+prompt = "91歳"
+generate_kwargs['prompt_ids'] = pipe.tokenizer.get_prompt_ids(prompt, return_tensors="pt").to(device)
+result = pipe(dataset[10]["audio"], generate_kwargs=generate_kwargs)
+result['text'] = result['text'][1 + len(prompt) + 1:]  # prompt has been added at the beginning of the output now, so remove it.
+print(result['text'])
+# あっぶったでもスルガさん、91歳、力強い走りに変わってきます。
 ```
 
 ### Additional Speed & Memory Improvements
@@ -310,31 +274,8 @@ pip install flash-attn --no-build-isolation
 Then pass `attn_implementation="flash_attention_2"` to `from_pretrained`:
 
 ```diff
-- model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
-+ model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, attn_implementation="flash_attention_2")
-```
-
-#### Torch Scale-Product-Attention (SDPA)
-
-If your GPU does not support Flash Attention, we recommend making use of PyTorch [scaled dot-product attention (SDPA)](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html). 
-This attention implementation is activated **by default** for PyTorch versions 2.1.1 or greater. To check 
-whether you have a compatible PyTorch version, run the following Python code snippet:
-
-```python
-from transformers.utils import is_torch_sdpa_available
-
-print(is_torch_sdpa_available())
-```
-
-If the above returns `True`, you have a valid version of PyTorch installed and SDPA is activated by default. If it 
-returns `False`, you need to upgrade your PyTorch version according to the [official instructions](https://pytorch.org/get-started/locally/)
-
-Once a valid PyTorch version is installed, SDPA is activated by default. It can also be set explicitly by specifying 
-`attn_implementation="sdpa"` as follows:
-
-```diff
-- model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
-+ model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, attn_implementation="sdpa")
+- model_kwargs = {"attn_implementation": "sdpa"} if torch.cuda.is_available() else {}
++ model_kwargs = {"attn_implementation": "flash_attention_2"} if torch.cuda.is_available() else {}
 ```