Views
No views yet
1import torch
2from qwen_asr import Qwen3ASRModel
3
4def get_dynamic_batches(items):
5 total = len(items)
6 i = 0
7
8 while i < total:
9 if items[i]['audio_duration_sec'] > 200:
10 batch_size = 1
11 else:
12 batch_size = int(5000 / items[i]['audio_duration_sec']) + 1
13 batch_size = min(batch_size, 64)
14
15 yield items[i:i + batch_size]
16 i += batch_size
17
18model = Qwen3ASRModel.from_pretrained(
19 "ZFTurbo/Qwen3-ASR-Children-Phonetic",
20 dtype=torch.bfloat16,
21 device_map="cuda:0",
22 max_inference_batch_size=64,
23 max_new_tokens=-1,
24)
25
26with torch.inference_mode():
27 for batch in get_dynamic_batches(items):
28 paths = []
29 languages = []
30 max_new_tokens = 0
31 for item in batch:
32 path = str(data_dir / item["audio_path"])
33 paths.append(path)
34 languages.append("English")
35 max_new_tokens = max(max_new_tokens, int(item["audio_duration_sec"] * 20))
36
37 print(
38 "Batch size:", len(batch),
39 "Duration:", batch[0]["audio_duration_sec"],
40 "Processed:", total,
41 "Max tokens:", max_new_tokens,
42 )
43 cur_time = time.time()
44
45 results = model.transcribe(
46 audio=paths,
47 language=languages, # can also be set to None for automatic language detection
48 return_time_stamps=False,
49 max_new_tokens=max_new_tokens,
50 )
51
52 predictions = {}
53 for i, r in enumerate(results):
54 predictions[i] = r.text
55