Update app.py

app.py

@@ -1,45 +1,219 @@
-import gradio as gr
-from transformers import pipeline
-import torch
+# -*- coding: utf-8 -*-
+"""demo 2/3.ipynb
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/drive/1QeNS57tZzvJudeNjQczKJ-PbN0l1tK6V
+
+# Import library
+"""
+import os
 import librosa
-import soundfile
+import gradio as gr
+import noisereduce as nr
+from scipy.io import wavfile
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
+
+"""# Load model"""
+
+from google.colab import drive
+import os
+drive.mount('/content/gdrive')
+
+# load model and processor"
+processor = WhisperProcessor.from_pretrained("/content/gdrive/MyDrive/ColabNotebookShared/Speech2TextHuyenNguyen/Model/FPTVinTest2")
+model = WhisperForConditionalGeneration.from_pretrained("/content/gdrive/MyDrive/ColabNotebookShared/Speech2TextHuyenNguyen/Model/FPTVinTest2/checkpoint-1332").to("cuda")
+model.config.forced_decoder_ids = processor.get_decoder_prompt_ids(task = "transcribe")
+
+"""# Slipt audio"""
+
+from pydub import AudioSegment
+
+def preprocessing(path):
+    # CONVERT MP3 -> WAV
+    type_file = path.split(".")[-1]
+    sound = AudioSegment.from_file(path, type_file)
+    path_list = []
+
+    # SPLIT AUDIO
+    time_audio = int(sound.duration_seconds / 20) + 1
+    for i in range(time_audio):
+        t1 = i * 20 * 1000
+        t2 = (i+1) * 20 * 1000
+        if i == (time_audio-1):
+          newAudio = sound[t1:]
+        else:
+          newAudio = sound[t1:t2]
+
+        newAudio = newAudio.split_to_mono()[0]
+        newAudio = newAudio.set_frame_rate(16000)    # convert frequency : mọi freq --> 16000kHz
+        audio_path = '/content/new_audio' + str(i) + '.wav'
+        newAudio.export(audio_path, format="wav") 
+        path_list.append(audio_path)
+    return path_list
+
+"""# Capitalization"""
+
+!git lfs install
+!git clone https://github.com/huyenxam/Vicap.git
+
+# Commented out IPython magic to ensure Python compatibility.
+# %cd {"/content/Vicap"}
+
+import os
+from gec_model import GecBERTModel
+
+cache_dir = "./"
+model_cap = GecBERTModel(
+    vocab_path=os.path.join(cache_dir, "vocabulary"),
+    model_paths="dragonSwing/vibert-capu",
+    split_chunk=True
+)
+
+"""# Spelling Correction"""
 
-SAMPLE_RATE = 16000
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+tokenizer_spell = AutoTokenizer.from_pretrained("VietAI/vit5-base")  
+model_spell = AutoModelForSeq2SeqLM.from_pretrained("HuyenNguyen/Vi-test1")
+model_spell.cuda()
 
-pipe = pipeline(model="HuyenNguyen/Vin-W-22000")
+def spelling_text(text):
+  encoding = tokenizer_spell(text, return_tensors="pt")
+  input_ids, attention_masks = encoding["input_ids"].to("cuda"), encoding["attention_mask"].to("cuda")
+  outputs = model_spell.generate(
+      input_ids=input_ids, attention_mask=attention_masks,
+      max_length=30,
+  )
 
+  for output in outputs:
+      line = tokenizer_spell.decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+  return line
 
-def transcribe(Microphone, File_Upload):
+def spelling(transcription):
+  sentences = transcription.split(" ")
+  len_sen = int(len(sentences) / 25) + 1
+  result = ""
+  for i in range(len_sen):
+      t1 = i * 24
+      t2 = (i+1) * 24 
+      if i == (len_sen - 1):
+        text = " ".join(sentences[t1:])
+      else:
+        text = " ".join(sentences[t1:t2])
+      result = result + " " + spelling_text(text)
+  
+  return result
+
+"""# Speech To Text"""
+
+import torch
+import numpy as np
+import gradio as gr
+from scipy.io.wavfile import write
+import pytube as pt
+from transformers import pipeline
+from huggingface_hub import model_info
+
+def transcribe(microphone, file_upload):
     warn_output = ""
-    if (Microphone is not None) and (File_Upload is not None):
-        warn_output = "WARNING: You've uploaded an audio file and used the microphone. " \
-                      "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
-        file = Microphone
+    if (microphone is not None) and (file_upload is not None):
+        warn_output = (
+            "WARNING: You've uploaded an audio file and used the microphone. "
+            "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
+        )
 
-    elif (Microphone is None) and (File_Upload is None):
+    elif (microphone is None) and (file_upload is None):
         return "ERROR: You have to either use the microphone or upload an audio file"
 
-    elif Microphone is not None:
-        file = Microphone
-    else:
-        file = File_Upload
+    path = microphone if microphone is not None else file_upload
+
+    X_new, sr_new = librosa.load(path)
+    dst = "/content/audio.wav"
+    write(dst, sr_new, X_new)
+
+    # Split audio
+    transcription = ""
+    path_list = preprocessing(dst)
+    for audio_path in path_list:
+      # X, sr = noise(audio_path)
+      X, sr = librosa.load(audio_path, sr=16000)
+      input_features = processor(X.astype('float16'), return_tensors="pt").input_features 
+      # predicted_ids = model.generate(input_features.to("cuda"), temperature=1.0)
+      predicted_ids = model.generate(input_features.to("cuda"))
+      text = processor.batch_decode(predicted_ids, skip_special_tokens = True)[0]
+      transcription = transcription + " " + text
+
+    transcription_spell = spelling(transcription)
     
-    text = pipe(file)["text"]
+    transcription_cap = model_cap(transcription_spell)[0]
+
+    # sentence_result = "Câu gốc: " + transcription + "\n" + "Câu sửa lỗi chính tả: " + transcription_spell + "\n" + "Thêm dấu: " + transcription_cap
+    return transcription_cap
+
 
-    return warn_output + text
+def _return_yt_html_embed(yt_url):
+    video_id = yt_url.split("?v=")[-1]
+    HTML_str = (
+        f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
+        " </center>"
+    )
+    return HTML_str
 
 
-iface = gr.Interface(
+def yt_transcribe(yt_url):
+    # yt = pt.YouTube(yt_url)
+    # html_embed_str = _return_yt_html_embed(yt_url)
+    # stream = yt.streams.filter(only_audio=True)[0]
+    # src = "/content/audio.mp3"
+    # dst = "/content/audio.wav"
+    # stream.download(filename=src)
+
+    # X_new, sr_new = librosa.load(src)
+
+    # write(dst, sr_new, X_new)
+    # # X_new, sr_new = librosa.load(src)
+    # path_list = preprocessing(dst)
+    # transcription = " "
+    # for audio_path in path_list:
+    #   # X, sr = noise(audio_path)
+    #   X, sr = librosa.load(audio_path, sr=16000)
+    #   input_features = processor(X.astype('float16'), return_tensors="pt").input_features 
+    #   predicted_ids = model.generate(input_features.to("cuda"))
+    #   text = processor.batch_decode(predicted_ids, skip_special_tokens = True)[0]
+    #   transcription = transcription + " " + text
+
+    # transcription = spelling(transcription)
+    # transcription = model_cap(transcription)[0]
+    return "ouput", 'This feature is temporarily locked'
+
+
+demo = gr.Blocks()
+
+mf_transcribe = gr.Interface(
     fn=transcribe,
     inputs=[
-        gr.inputs.Audio(source="microphone", type='filepath', optional=True),
-        gr.inputs.Audio(source="upload", type='filepath', optional=True),
+        gr.inputs.Audio(source="microphone", type="filepath", optional=True),
+        gr.inputs.Audio(source="upload", type="filepath", optional=True),
     ],
     outputs="text",
     layout="horizontal",
     theme="huggingface",
-    title="Speed to Text Huyen Nguyen",
-    allow_flagging='never',
+    title="PYLAB Demo: Transcribe Audio",
+    allow_flagging="never",
 )
 
-iface.launch(enable_queue=True)
+yt_transcribe = gr.Interface(
+    fn=yt_transcribe,
+    inputs=[gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL")],
+    outputs=["html", "text"],
+    layout="horizontal",
+    theme="huggingface",
+    title="PYLAB Demo: Transcribe YouTube",
+    allow_flagging="never",
+)
+
+with demo:
+    gr.TabbedInterface([mf_transcribe, yt_transcribe], ["Transcribe Audio", "Transcribe YouTube"])
+
+demo.launch(enable_queue=True)