-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
105 lines (85 loc) · 3.61 KB
/
main.py
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
import io
import os
import torch
import pandas as pd
import urllib
import tarfile
import whisper
from scipy.io import wavfile
from tqdm import tqdm
class Fleurs(torch.utils.data.Dataset):
"""
A simple class to wrap Fleurs and subsample a portion of the dataset as needed.
"""
def __init__(self, lang, device, split="test", subsample_rate=1):
url = f"https://storage.googleapis.com/xtreme_translations/FLEURS102/{lang}.tar.gz"
tar_path = os.path.expanduser(f"~/.cache/fleurs/{lang}.tgz")
os.makedirs(os.path.dirname(tar_path), exist_ok=True)
if not os.path.exists(tar_path):
with urllib.request.urlopen(url) as source, open(tar_path, "wb") as output:
with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True,
unit_divisor=1024) as loop:
while True:
buffer = source.read(8192)
if not buffer:
break
output.write(buffer)
loop.update(len(buffer))
labels = {}
all_audio = {}
with tarfile.open(tar_path, "r:gz") as tar:
for member in tar.getmembers():
name = member.name
if name.endswith(f"{split}.tsv"):
labels = pd.read_table(tar.extractfile(member), names=(
"id", "file_name", "raw_transcription", "transcription", "_", "num_samples", "gender"))
if f"/{split}/" in name and name.endswith(".wav"):
audio_bytes = tar.extractfile(member).read()
all_audio[os.path.basename(name)] = wavfile.read(io.BytesIO(audio_bytes))[1]
self.labels = labels.to_dict("records")[::subsample_rate]
self.all_audio = all_audio
self.device = device
def __len__(self):
return len(self.labels)
def __getitem__(self, item):
record = self.labels[item]
audio = torch.from_numpy(self.all_audio[record["file_name"]].copy())
text = record["transcription"]
return (audio, text)
# Display options for pandas dataset
pd.options.display.max_rows = 100
pd.options.display.max_colwidth = 1000
# Set inference device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Set device to {device}")
# Set language (korean)
language_google = "ko_kr"
language_whisper = "korean"
print(f"Set language to {language_whisper}")
# Create dataset object, selecting only 10 examples for brevity
dataset = Fleurs(language_google, subsample_rate=1, device=device)
dataset = torch.utils.data.random_split(dataset, [10, len(dataset)-10])[0]
print(f"Created dataset")
# Load tiny Whisper model
model_size = "tiny"
model = whisper.load_model(model_size)
print(f"Loaded {model_size} Whisper model")
# Set options
options = dict(language=language_whisper, beam_size=5, best_of=5)
transcribe_options = dict(task="transcribe", **options)
translate_options = dict(task="translate", **options)
# Run inference
references = []
transcriptions = []
translations = []
for idx, (audio, text) in enumerate(tqdm(dataset)):
print(f"\nProcessing datum number {idx+1}")
transcription = model.transcribe(audio, **transcribe_options)["text"]
translation = model.transcribe(audio, **translate_options)["text"]
transcriptions.append(transcription)
translations.append(translation)
references.append(text)
# Create dataframe from results and save the data
data = pd.DataFrame(dict(reference=references, transcription=transcriptions, translation=translations))
print(data)
data.to_csv("results.csv")