FilesExpand file tree

master

/

bench.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

225 lines (183 loc) · 6.74 KB

master

/

bench.py

Copy path

Top

File metadata and controls

• Code
• Blame

225 lines (183 loc) · 6.74 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

import os

import subprocess

import re

import csv

import wave

import contextlib

import argparse

# Custom action to handle comma-separated list

class ListAction(argparse.Action):

def __call__(self, parser, namespace, values, option_string=None):

setattr(namespace, self.dest, [int(val) for val in values.split(",")])

parser = argparse.ArgumentParser(description="Benchmark the speech recognition model")

# Define the argument to accept a list

parser.add_argument(

"-t",

"--threads",

dest="threads",

action=ListAction,

default=[4],

help="List of thread counts to benchmark (comma-separated, default: 4)",

)

parser.add_argument(

"-p",

"--processors",

dest="processors",

action=ListAction,

default=[1],

help="List of processor counts to benchmark (comma-separated, default: 1)",

)

parser.add_argument(

"-f",

"--filename",

type=str,

default="./samples/jfk.wav",

help="Relative path of the file to transcribe (default: ./samples/jfk.wav)",

)

# Parse the command line arguments

args = parser.parse_args()

sample_file = args.filename

threads = args.threads

processors = args.processors

# Define the models, threads, and processor counts to benchmark

models = [

"ggml-tiny.en.bin",

"ggml-tiny.bin",

"ggml-base.en.bin",

"ggml-base.bin",

"ggml-small.en.bin",

"ggml-small.bin",

"ggml-medium.en.bin",

"ggml-medium.bin",

"ggml-large-v1.bin",

"ggml-large-v2.bin",

"ggml-large-v3.bin",

"ggml-large-v3-turbo.bin",

]

metal_device = ""

# Initialize a dictionary to hold the results

results = {}

gitHashHeader = "Commit"

modelHeader = "Model"

hardwareHeader = "Hardware"

recordingLengthHeader = "Recording Length (seconds)"

threadHeader = "Thread"

processorCountHeader = "Processor Count"

loadTimeHeader = "Load Time (ms)"

sampleTimeHeader = "Sample Time (ms)"

encodeTimeHeader = "Encode Time (ms)"

decodeTimeHeader = "Decode Time (ms)"

sampleTimePerRunHeader = "Sample Time per Run (ms)"

encodeTimePerRunHeader = "Encode Time per Run (ms)"

decodeTimePerRunHeader = "Decode Time per Run (ms)"

totalTimeHeader = "Total Time (ms)"

def check_file_exists(file: str) -> bool:

return os.path.isfile(file)

def get_git_short_hash() -> str:

try:

return (

subprocess.check_output(["git", "rev-parse", "--short", "HEAD"])

.decode()

.strip()

)

except subprocess.CalledProcessError as e:

return ""

def wav_file_length(file: str = sample_file) -> float:

with contextlib.closing(wave.open(file, "r")) as f:

frames = f.getnframes()

rate = f.getframerate()

duration = frames / float(rate)

return duration

def extract_metrics(output: str, label: str) -> tuple[float, float]:

match = re.search(rf"{label} \s*=\s*(\d+\.\d+)\s*ms\s*/\s*(\d+)\s*runs", output)

time = float(match.group(1)) if match else None

runs = float(match.group(2)) if match else None

return time, runs

def extract_device(output: str) -> str:

match = re.search(r"picking default device: (.*)", output)

device = match.group(1) if match else "Not found"

return device

# Check if the sample file exists

if not check_file_exists(sample_file):

raise FileNotFoundError(f"Sample file {sample_file} not found")

recording_length = wav_file_length()

# Check that all models exist

# Filter out models from list that are not downloaded

filtered_models = []

for model in models:

if check_file_exists(f"models/{model}"):

filtered_models.append(model)

else:

print(f"Model {model} not found, removing from list")

models = filtered_models

# Loop over each combination of parameters

for model in filtered_models:

for thread in threads:

for processor_count in processors:

# Construct the command to run

cmd = f"./build/bin/whisper-cli -m models/{model} -t {thread} -p {processor_count} -f {sample_file}"

# Run the command and get the output

process = subprocess.Popen(

cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT

)

output = ""

while process.poll() is None:

output += process.stdout.read().decode()

# Parse the output

load_time_match = re.search(r"load time\s*=\s*(\d+\.\d+)\s*ms", output)

load_time = float(load_time_match.group(1)) if load_time_match else None

metal_device = extract_device(output)

sample_time, sample_runs = extract_metrics(output, "sample time")

encode_time, encode_runs = extract_metrics(output, "encode time")

decode_time, decode_runs = extract_metrics(output, "decode time")

total_time_match = re.search(r"total time\s*=\s*(\d+\.\d+)\s*ms", output)

total_time = float(total_time_match.group(1)) if total_time_match else None

model_name = model.replace("ggml-", "").replace(".bin", "")

print(

f"Ran model={model_name} threads={thread} processor_count={processor_count}, took {total_time}ms"

)

# Store the times in the results dictionary

results[(model_name, thread, processor_count)] = {

loadTimeHeader: load_time,

sampleTimeHeader: sample_time,

encodeTimeHeader: encode_time,

decodeTimeHeader: decode_time,

sampleTimePerRunHeader: round(sample_time / sample_runs, 2),

encodeTimePerRunHeader: round(encode_time / encode_runs, 2),

decodeTimePerRunHeader: round(decode_time / decode_runs, 2),

totalTimeHeader: total_time,

}

# Write the results to a CSV file

with open("benchmark_results.csv", "w", newline="") as csvfile:

fieldnames = [

gitHashHeader,

modelHeader,

hardwareHeader,

recordingLengthHeader,

threadHeader,

processorCountHeader,

loadTimeHeader,

sampleTimeHeader,

encodeTimeHeader,

decodeTimeHeader,

sampleTimePerRunHeader,

encodeTimePerRunHeader,

decodeTimePerRunHeader,

totalTimeHeader,

]

writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

writer.writeheader()

shortHash = get_git_short_hash()

# Sort the results by total time in ascending order

sorted_results = sorted(results.items(), key=lambda x: x[1].get(totalTimeHeader, 0))

for params, times in sorted_results:

row = {

gitHashHeader: shortHash,

modelHeader: params[0],

hardwareHeader: metal_device,

recordingLengthHeader: recording_length,

threadHeader: params[1],

processorCountHeader: params[2],

}

row.update(times)

writer.writerow(row)