Update README.md

README.md

@@ -80,6 +80,8 @@ python3 -m pip install transformers@git+https://github.com/huggingface/transform
 ```
 as this is the transformers version we are using when building this model.
 
+### Simple Demo
+
 ```python
 from transformers import AutoProcessor, AutoModelForCausalLM
 
@@ -125,6 +127,58 @@ cont = outputs[:, inputs["input_ids"].shape[-1] :]
 print(processor.batch_decode(cont, skip_special_tokens=True)[0])
 ```
 
+### Batch Inference
+The model supports batch inference with transformers. An example demo is like this:
+```python
+from transformers import AutoProcessor, AutoModelForCausalLM
+
+import torch
+import librosa
+
+def load_audio():
+    return librosa.load(librosa.ex("libri1"), sr=16000)[0]
+
+def load_audio_2():
+    return librosa.load(librosa.ex("libri2"), sr=16000)[0]
+
+
+processor = AutoProcessor.from_pretrained("lmms-lab/Aero-1-Audio-1.5B", trust_remote_code=True)
+# We encourage to use flash attention 2 for better performance
+# Please install it with `pip install --no-build-isolation flash-attn`
+# If you do not want flash attn, please use sdpa or eager`
+model = AutoModelForCausalLM.from_pretrained("lmms-lab/Aero-1-Audio-1.5B", device_map="cuda", torch_dtype="auto", attn_implementation="flash_attention_2", trust_remote_code=True)
+model.eval()
+
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "audio_url",
+                "audio": "placeholder",
+            },
+            {
+                "type": "text",
+                "text": "Please transcribe the audio",
+            }
+        ]
+    }
+]
+messages = [messages, messages]
+
+audios = [load_audio(), load_audio_2()]
+
+processor.tokenizer.padding_side="left"
+prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+inputs = processor(text=prompt, audios=audios, sampling_rate=16000, return_tensors="pt", padding=True)
+inputs = {k: v.to("cuda") for k, v in inputs.items()}
+outputs = model.generate(**inputs, eos_token_id=151645, pad_token_id=151643, max_new_tokens=4096)
+
+cont = outputs[:, inputs["input_ids"].shape[-1] :]
+
+print(processor.batch_decode(cont, skip_special_tokens=True))
+```
+
 ## Training Details
 
 ### Training Data