Upload core.py

core.py

@@ -0,0 +1,78 @@
+import transformers 
+from transformers import pipeline
+
+import whisper
+
+import datetime
+
+transformers.utils.move_cache()
+
+# ====================================
+# Load speech recognition model
+# speech_recognition_pipeline = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h")
+speech_recognition_model = whisper.load_model("base")
+
+# ====================================
+# Load text summarization model English
+# text_summarization_pipeline_En = pipeline("summarization", model="facebook/bart-large-cnn")
+tokenizer_En = transformers.AutoTokenizer.from_pretrained("facebook/bart-large-cnn")  
+text_summarization_model_En = transformers.AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
+
+# ====================================
+# Load text summarization model Vietnamese
+tokenizer_Vi = transformers.AutoTokenizer.from_pretrained("VietAI/vit5-large-vietnews-summarization")  
+text_summarization_model_Vi = transformers.AutoModelForSeq2SeqLM.from_pretrained("VietAI/vit5-large-vietnews-summarization")
+
+def asr_transcript(input_file):
+    audio = whisper.load_audio(input_file)
+    output = speech_recognition_model.transcribe(audio)
+    text = output['text']
+    lang = "English"
+    if output["language"] == 'en':
+        lang = "English"
+    elif output["language"] == 'vi':
+        lang = "Vietnamese"
+
+    detail = ""
+    for segment in output['segments']:
+        start = str(datetime.timedelta(seconds=round(segment['start'])))
+        end = str(datetime.timedelta(seconds=round(segment['end'])))
+        small_text = segment['text']
+        detail = detail + start + "-" + end + " " + small_text + "\n"
+    return text, lang, detail
+
+def text_summarize_en(text_input):
+    encoding = tokenizer_En(text_input, truncation=True, return_tensors="pt")
+    input_ids, attention_masks = encoding["input_ids"], encoding["attention_mask"]
+    outputs = text_summarization_model_En.generate(
+        input_ids=input_ids, attention_mask=attention_masks,
+        max_length=256,
+        early_stopping=True
+    )
+    text = ""
+    for output in outputs:
+        line = tokenizer_En.decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+        text = text + line
+    return text
+
+def text_summarize_vi(text_input):
+    encoding = tokenizer_Vi(text_input, truncation=True, return_tensors="pt")
+    input_ids, attention_masks = encoding["input_ids"], encoding["attention_mask"]
+    outputs = text_summarization_model_Vi.generate(
+        input_ids=input_ids, attention_mask=attention_masks,
+        max_length=256,
+        early_stopping=True
+    )
+    text = ""
+    for output in outputs:
+        line = tokenizer_Vi.decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+        text = text + line
+    return text
+
+def text_summarize(text_input, lang):
+    if lang == 'English':
+        return text_summarize_en(text_input)
+    elif lang == 'Vietnamese':
+        return text_summarize_vi(text_input)
+    else:
+        return ""