sfedar commited on
Commit
6b64d77
·
verified ·
1 Parent(s): 799bc28

Add a separate pipeline for translation (Helsinki-NLP/opus-mt model)

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -8,8 +8,9 @@ from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Proce
8
 
9
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
10
 
11
- # load speech translation checkpoint
12
  asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
 
13
 
14
  # load text-to-speech checkpoint and speaker embeddings
15
  tts_model_name = "sanchit-gandhi/speecht5_tts_vox_nl"
@@ -22,8 +23,9 @@ speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze
22
 
23
 
24
  def translate(audio):
25
- outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate", "language": "dutch"})
26
- return outputs["text"]
 
27
 
28
 
29
  def synthesise(text):
 
8
 
9
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
10
 
11
+ # load speech translation checkpoints
12
  asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
13
+ translation_pipeline = pipeline("translation", model="Helsinki-NLP/opus-mt-en-nl", device=device)
14
 
15
  # load text-to-speech checkpoint and speaker embeddings
16
  tts_model_name = "sanchit-gandhi/speecht5_tts_vox_nl"
 
23
 
24
 
25
  def translate(audio):
26
+ transcripts = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})["text"]
27
+ outputs = translation_pipeline(transcripts)
28
+ return outputs[0]['translation_text']
29
 
30
 
31
  def synthesise(text):