TuringVideos / app.py
EnzoBustos's picture
Update app.py
2cf872d
raw
history blame
3.95 kB
import yt_dlp
import whisper
import os
import streamlit as st
import transformers
from transformers import pipeline
from summarizer import Summarizer
import torch
import ffmpeg
import speech_recognition as sr
from moviepy.editor import AudioFileClip
#Download youtube video
def download_audio(link):
with yt_dlp.YoutubeDL({'extract_audio': True, 'format': 'bestaudio', 'outtmpl': 'video.mp3'}) as video:
video.download(link)
def get_transcription_en(mp3_file):
model = whisper.load_model("tiny.en")
directory = os.getcwd()
result = model.transcribe(os.path.join(directory, mp3_file))
return result["text"]
#def portuguese_sum_pipeline(link):
# video_title = download_audio(link)
def english_sum_pipeline(link):
download_audio(link)
transcript_text = get_transcription_en("video.mp3")
#extractive summarization
extractive_model = Summarizer()
extractive = extractive_model(transcript_text, num_sentences=15)
#abstractive summarization
device_num = 0 if torch.cuda.is_available() else -1
abstractive_summarizer = pipeline("summarization", model="facebook/bart-large-cnn", tokenizer="facebook/bart-large-cnn", device=device_num)
output_text = abstractive_summarizer(extractive)[0]["summary_text"]
return transcript_text, output_text
def english_qa_pipeline(question, context):
nlp = pipeline("question-answering", model='distilbert-base-uncased-distilled-squad')
result = nlp(question=question, context=context)
return result["answer"]
#Collect inputs and create the interface
def main():
header = st.container()
model = st.container()
model_1, model_2 = st.columns(2)
qa = st.container()
qa_1, qa_2 = st.columns(2)
with header:
st.title("TuringVideos")
st.write("Este trabalho visa a criação de uma interface capaz de sumarizar e responder perguntas sobre um determinado vídeo em português ou inglês!")
with model:
st.header("Modelo para sumarização")
with model_1:
language = st.selectbox('Qual a linguagem do seu modelo?', ('Português (pt)', 'Inglês (en)', 'Outra'))
link = st.text_area(label="Coloque o link do seu vídeo do YouTube!", height=25, placeholder="Digite seu link...")
submit_1 = st.button('Gerar sumário!')
with model_2:
if submit_1:
with st.spinner('Wait for it...'):
if language == 'Português (pt)':
#outputs = portuguese_sum_pipeline(link)
st.write("Modelo ainda não implementado.")
elif language == 'Inglês (en)':
outputs = english_sum_pipeline(link)
else:
st.write("Erro na seleção de linguagem.")
st.write("Sumário.....................................................................: \n {} \n \n".format(outputs[1]))
with qa:
st.header("Modelo para resposta de perguntas")
with qa_1:
question = st.text_area(label="Digite sua pergunta para o modelo", height=25, placeholder="Digite sua pergunta...")
submit_2 = st.button('Gerar resposta!')
with qa_2:
with st.spinner('Wait for it...'):
if submit_2:
if language == 'Português (pt)':
#answer = portuguese_qa_pipeline(link)
st.write("Modelo ainda não implementado.")
elif language == 'Inglês (en)':
answer = english_qa_pipeline(question, outputs[0])
else:
st.write("Erro na seleção de linguagem.")
st.write("Resposta.....................................................................: \n {} \n \n".format(answer))
main()