9farccontioshi/TranscriptApi

1import os2import librosa3import soundfile as sf4from pytube import YouTube5import urllib.parse as urlparse6from moviepy.editor import VideoFileClip7import shutil8import whisper9import torch 10from transformers import pipeline11from tqdm.auto import tqdm12from PyPDF2 import PdfReader13 14 15device = 'cuda' if torch.cuda.is_available() else 'cpu'16# device = 'cpu'17 18 19checkpoint = 'Th3BossC/SummarizationModel_t5-small_opeai_tldr'20 21 22 23 24 25 26############### video queries ###############27def title(video_id):28    return YouTube('https://www.youtube.com/watch?v=' + video_id).title29 30def get_video_id(video_url):31    url_data = urlparse.urlparse("http://www.youtube.com/watch?v=z_AbfPXTKms&NR=1")32    query = urlparse.parse_qs(url_data.query)33    video = query["v"][0]34    return video35 36def get_video(video_url, location, filename = 'audio'):37    if not os.path.exists(location):38        os.makedirs(location)39    video_filename = location + filename + '.mp4'40    audio_filename = location + filename + '.mp3'41    print('[INFO] downloading video...')42    video = YouTube(video_url).streams.filter(file_extension = 'mp4').first().download(filename = video_filename)43    video = VideoFileClip(video_filename)44    print('[INFO] extracting audio from video...')45    video.audio.write_audiofile(audio_filename)46    #os.remove(video_filename)47 48    return audio_filename49 50############################################################51 52 53############### Audio ###############54def chunk_audio(filename, segment_length, output_dir):55    if not os.path.isdir(output_dir):56        os.mkdir(output_dir)57    audio, sr = librosa.load(filename, sr = 44100)58    duration = librosa.get_duration(y = audio, sr = sr)59    num_segments = int(duration / segment_length) + 160    print(f'[INFO] Chunking {num_segments} chunks...')61 62    audio_files = []63 64    for i in range(num_segments):65        start = i*segment_length*sr 66        end = (i+1)*segment_length*sr67        segment = audio[start:end]68        sf.write(os.path.join(output_dir, f"segment_{i}.mp3"), segment, sr)69        audio_files.append(output_dir + f'segment_{i}.mp3')70        71    print(audio_files)72    #os.remove(filename)73    return audio_files74 75def transcribe_audio(audio_files, output_file = None, model = whisper.load_model('base', device = device)):76    print('[INFO] converting audio to text...')77    transcripts = []78    model.to(device)79    for audio_file in audio_files:80        response = model.transcribe(audio_file)81        transcripts.append(response['text'])82 83    if output_file is not None:84        with open(output_file, 'w') as f:85            for transcript in transcripts:86                f.write(transcript + '\n')87    88    return transcripts89 90############################################################91 92 93############################################################94    95############### Compile all functions ###############96def summarize_youtube_video(video_url, outputs_dir):97    print(f'[INFO] running on {device}')98    raw_audio_dir = f'{outputs_dir}/raw_audio/'99    chunks_dir = f'{outputs_dir}/chunks/'100    transcripts_file = f'{outputs_dir}/transcripts.txt'101    summary_file = f'{outputs_dir}/summary.txt'102    segment_length = 60*10103 104    if os.path.exists(outputs_dir):105        shutil.rmtree(outputs_dir)106        os.mkdir(outputs_dir)107 108    audio_filename = get_video(video_url, raw_audio_dir)109    chunked_audio_files = chunk_audio(audio_filename, segment_length, chunks_dir)110    transcriptions = transcribe_audio(chunked_audio_files, transcripts_file)111 112 113    # splitting transcription into sentences    114    sentences = [] 115    for transcript in transcriptions:116        sentences += transcript.split('.')117 118    sentences_len = [len(sentence) for sentence in sentences]119    sentence_mean_length = sum(sentences_len) // len(sentences_len)120 121    num_sentences_per_step = int(1600 / (sentence_mean_length))122    num_steps = (len(sentences) // num_sentences_per_step) + (len(sentences) % num_sentences_per_step != 0)123 124    print(f"""125    [INFO] sentences_len : {len(sentences_len)}126    [INFO] sentence_mean_length : {sentence_mean_length},127    [INFO] num_sentences_per_step : {num_sentences_per_step},128    [INFO] num_steps : {num_steps}129    """)130 131    summarizer = pipeline('summarization', model = checkpoint, tokenizer = checkpoint, max_length = 200, truncation = True)132 133    summaries = []134 135    for i in tqdm(range(num_steps)):136        chunk = ' '.join(sentences[num_sentences_per_step*i : num_sentences_per_step*(i+1)])137        summary = summarizer(chunk, do_sample = False)[0]['summary_text']138        summaries.append(summary)139 140    complete_summary = ' '.join(summaries)141    with open(summary_file, 'w') as f:142        f.write(complete_summary)143 144    with open(transcripts_file, 'r') as f:145        complete_transcript = f.read()146    return {'transcript': complete_transcript, 'summary' : complete_summary}147############################################################148 149 150 151############ File Summarize ############152 153def extract_text_pdf(file_location = 'TranscriptApi/static/files/temp.pdf'):154    reader = PdfReader(file_location)155    text = ""156    for page in reader.pages:157        text += page.extract_text()158    return text;159 160def extract_text_txt(file_location = 'TranscriptApi/static/files/temp.txt'):161    with open(file_location, "r") as f:162        text = f.read()163    return text164 165 166 167 168def summarize_string(text : str):169    sentences = text.split('.')170 171    summarizer = pipeline('summarization', model = checkpoint, tokenizer = checkpoint, max_length = 200, truncation = True, device = 0)172 173    sentences_len = [len(sentence) for sentence in sentences]174    sentence_mean_length = sum(sentences_len) // len(sentences_len)175 176    num_sentences_per_step = int(1600 / (sentence_mean_length))177    num_steps = (len(sentences) // num_sentences_per_step) + (len(sentences) % num_sentences_per_step != 0)178 179    print(f"""180    [INFO] sentences_len : {len(sentences_len)}181    [INFO] sentence_mean_length : {sentence_mean_length},182    [INFO] num_sentences_per_step : {num_sentences_per_step},183    [INFO] num_steps : {num_steps}184    """)185 186 187    summaries = []188    for i in tqdm(range(num_steps)):189        chunk = ' '.join(sentences[num_sentences_per_step*i : num_sentences_per_step*(i+1)])190        summary = summarizer(chunk, do_sample = False)[0]['summary_text']191        summaries.append(summary)192 193    complete_summary = ' '.join(summaries)194    return complete_summary195 196 197################################################198 199 200def summarize_file(file_location, file_extension, working_dir = "TranscriptApi/static/files"):201    # _, file_extension = os.path.splitext(file_location)202    text = ""203    if file_extension == 'pdf':204        text = extract_text_pdf(file_location)205    elif file_extension == 'txt':206        text = extract_text_txt(file_location)207    else:208        return "[ERROR]"209    210    if os.path.exists(working_dir):211        shutil.rmtree(working_dir)212    os.mkdir(working_dir)213    return [text, summarize_string(text)]214 215def answer(question: str, context : str):216    # qa = pipeline(task = "question-answering", model = "Th3BossC/QuestionAnsweringModel", tokenizer = "Th3BossC/QuestionAnsweringModel")217    qa = pipeline(task = "question-answering", model = "deepset/roberta-base-squad2")218    return qa(question = question, context = context)['answer']