Pooya-Fallah commited on
Commit
9dc2477
·
verified ·
1 Parent(s): 464e4ce

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +75 -0
  2. requirements.txt +12 -0
app.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import nemo.collections.asr as nemo_asr
3
+ import gc
4
+ import numpy as np
5
+ import torchaudio
6
+
7
+ pretrained_model_path="./stt_fa_fastconformer_hybrid_large_finetuned.nemo"
8
+
9
+ # Clear up memory
10
+ torch.cuda.empty_cache()
11
+ gc.collect()
12
+ model = nemo_asr.models.EncDecHybridRNNTCTCModel.restore_from(pretrained_model_path)
13
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
14
+ # device = 'cpu' # You can transcribe even longer samples on the CPU, though it will take much longer !
15
+ model = model.to(device)
16
+ model.freeze()
17
+
18
+ def transcribe(stream, new_chunk):
19
+ if new_chunk is None:
20
+ return None, ""
21
+ # 'audio' is a tuple: (sample_rate, data)
22
+ sample_rate, data = new_chunk
23
+
24
+ # Ensure the model is on the correct device
25
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
26
+
27
+ # Convert audio data to the expected format
28
+ if isinstance(data, np.ndarray):
29
+ audio_tensor = torch.tensor(data, dtype=torch.float32)
30
+ else:
31
+ raise ValueError("Audio data must be a numpy array")
32
+
33
+ # Resample if sample rate is not 16000
34
+ target_sample_rate = 16000
35
+ if sample_rate != target_sample_rate:
36
+ resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sample_rate)
37
+ audio_tensor = resampler(audio_tensor)
38
+
39
+ if stream is not None:
40
+ stream['audio'] = torch.cat([stream['audio'], audio_tensor], dim=-1)
41
+ else:
42
+ stream = {"text": ""}
43
+ stream['audio'] = audio_tensor
44
+
45
+
46
+ max_length = 5 * target_sample_rate # 5 seconds
47
+ new_text = ""
48
+
49
+ # Process all chunks that fit max_length
50
+ while stream['audio'].shape[-1] > max_length:
51
+ # Extract first max_length samples
52
+ audio_chunk = stream['audio'][..., :max_length]
53
+
54
+ # Transcribe
55
+ with torch.no_grad():
56
+ transcript = model.transcribe(audio_chunk) # Add batch dimension if needed
57
+
58
+ # Update text (adjust based on model's output format)
59
+ new_text += " " + transcript[0][0].strip() # Example adjustment
60
+
61
+ # Remove processed audio from buffer
62
+ stream['audio'] = stream['audio'][..., max_length:]
63
+
64
+ stream['text'] += new_text
65
+ return stream, stream['text'].strip()
66
+
67
+
68
+ interface = gr.Interface(
69
+ fn=transcribe,
70
+ inputs=['state', gr.Audio(sources="microphone", streaming=True, type="numpy")],
71
+ outputs=["state", "text"],
72
+ live=True,
73
+ )
74
+
75
+ interface.launch()
requirements.txt ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Python packages
2
+ wget
3
+ text-unidecode
4
+ matplotlib>=3.3.2
5
+ ffmpeg-python
6
+ gradio
7
+ numpy
8
+ torch
9
+ torchaudio
10
+
11
+ # Install NeMo from the Git repository (branch: main)
12
+ git+https://github.com/NVIDIA/NeMo.git@main#egg=nemo_toolkit[all]