from shared import DatasetArguments from utils import jaccard from functools import lru_cache from datetime import datetime import itertools from typing import Optional import model as model_module import segment from tqdm import tqdm from dataclasses import dataclass, field from transformers import HfArgumentParser from shared import extract_sponsor_matches_from_text, ACTION_OPTIONS, CATEGORIES, CATGEGORY_OPTIONS, START_SEGMENT_TEMPLATE, END_SEGMENT_TEMPLATE, GeneralArguments, CustomTokens import csv import re import random import logging from youtube_transcript_api import YouTubeTranscriptApi, CouldNotRetrieveTranscript, YouTubeRequestFailed, TooManyRequests import os import json import time import requests logging.basicConfig() logger = logging.getLogger(__name__) PROFANITY_RAW = '[ __ ]' # How YouTube transcribes profanity PROFANITY_CONVERTED = '*****' # Safer version for tokenizing NUM_DECIMALS = 3 # https://www.fincher.org/Utilities/CountryLanguageList.shtml # https://lingohub.com/developers/supported-locales/language-designators-with-regions LANGUAGE_PREFERENCE_LIST = ['en-GB', 'en-US', 'en-CA', 'en-AU', 'en-NZ', 'en-ZA', 'en-IE', 'en-IN', 'en-JM', 'en-BZ', 'en-TT', 'en-PH', 'en-ZW', 'en'] def parse_transcript_json(json_data, granularity): assert json_data['wireMagic'] == 'pb3' assert granularity in ('word', 'chunk') # TODO remove bracketed words? # (kiss smacks) # (upbeat music) # [text goes here] # Some manual transcripts aren't that well formatted... but do have punctuation # https://www.youtube.com/watch?v=LR9FtWVjk2c parsed_transcript = [] events = json_data['events'] for event_index, event in enumerate(events): segments = event.get('segs') if not segments: continue # This value is known (when phrase appears on screen) start_ms = event['tStartMs'] total_characters = 0 new_segments = [] for seg in segments: # Replace \n, \t, etc. with space text = ' '.join(seg['utf8'].split()) # Remove zero-width spaces and strip trailing and leading whitespace text = text.replace('\u200b', '').replace('\u200c', '').replace( '\u200d', '').replace('\ufeff', '').strip() # Alternatively, # text = text.encode('ascii', 'ignore').decode() # Needed for auto-generated transcripts text = text.replace(PROFANITY_RAW, PROFANITY_CONVERTED) if not text: continue offset_ms = seg.get('tOffsetMs', 0) new_segments.append({ 'text': text, 'start': round((start_ms + offset_ms)/1000, NUM_DECIMALS) }) total_characters += len(text) if not new_segments: continue if event_index < len(events) - 1: next_start_ms = events[event_index + 1]['tStartMs'] total_event_duration_ms = min( event.get('dDurationMs', float('inf')), next_start_ms - start_ms) else: total_event_duration_ms = event.get('dDurationMs', 0) # Ensure duration is non-negative total_event_duration_ms = max(total_event_duration_ms, 0) avg_seconds_per_character = ( total_event_duration_ms/total_characters)/1000 num_char_count = 0 for seg_index, seg in enumerate(new_segments): num_char_count += len(seg['text']) # Estimate segment end seg_end = seg['start'] + \ (num_char_count * avg_seconds_per_character) if seg_index < len(new_segments) - 1: # Do not allow longer than next seg_end = min(seg_end, new_segments[seg_index+1]['start']) seg['end'] = round(seg_end, NUM_DECIMALS) parsed_transcript.append(seg) final_parsed_transcript = [] for i in range(len(parsed_transcript)): word_level = granularity == 'word' if word_level: split_text = parsed_transcript[i]['text'].split() elif granularity == 'chunk': # Split on space after punctuation split_text = re.split( r'(?<=[.!?,-;])\s+', parsed_transcript[i]['text']) if len(split_text) == 1: split_on_whitespace = parsed_transcript[i]['text'].split() if len(split_on_whitespace) >= 8: # Too many words # Rather split on whitespace instead of punctuation split_text = split_on_whitespace else: word_level = True else: raise ValueError('Unknown granularity') segment_end = parsed_transcript[i]['end'] if i < len(parsed_transcript) - 1: segment_end = min(segment_end, parsed_transcript[i+1]['start']) segment_duration = segment_end - parsed_transcript[i]['start'] num_chars_in_text = sum(map(len, split_text)) num_char_count = 0 current_offset = 0 for s in split_text: num_char_count += len(s) next_offset = (num_char_count/num_chars_in_text) * segment_duration word_start = round( parsed_transcript[i]['start'] + current_offset, NUM_DECIMALS) word_end = round( parsed_transcript[i]['start'] + next_offset, NUM_DECIMALS) # Make the reasonable assumption that min wps is 1.5 final_parsed_transcript.append({ 'text': s, 'start': word_start, 'end': min(word_end, word_start + 1.5) if word_level else word_end }) current_offset = next_offset return final_parsed_transcript def list_transcripts(video_id): try: return YouTubeTranscriptApi.list_transcripts(video_id) except json.decoder.JSONDecodeError: return None WORDS_TO_REMOVE = [ CustomTokens.MUSIC.value, CustomTokens.APPLAUSE.value, CustomTokens.LAUGHTER.value ] @lru_cache(maxsize=16) def get_words(video_id, process=True, transcript_type='auto', fallback='manual', filter_words_to_remove=True, download=False, granularity='word'): """Get parsed video transcript with caching system returns None if not processed yet and process is False """ # NOTE: granularity='chunk' should only be used for generating training data... nowhere else transcript_path = os.path.join( # TODO use relative path to this 'transcripts', transcript_type, f'{video_id}.json') raw_transcript_json = None try: if not download and os.path.exists(transcript_path): # Load from file with open(transcript_path) as fp: raw_transcript_json = json.load(fp) # May be empty elif process: transcript_list = list_transcripts(video_id) if transcript_list is not None: if transcript_type == 'manual': ts = transcript_list.find_manually_created_transcript( LANGUAGE_PREFERENCE_LIST) else: ts = transcript_list.find_generated_transcript( LANGUAGE_PREFERENCE_LIST) raw_transcript = ts._http_client.get( f'{ts._url}&fmt=json3').content if raw_transcript: raw_transcript_json = json.loads(raw_transcript) except (TooManyRequests, YouTubeRequestFailed): raise # Cannot recover from these errors and do not mark as empty transcript except requests.exceptions.RequestException: # Can recover time.sleep(10) # Timeout return get_words(video_id, process, transcript_type, fallback, granularity) except CouldNotRetrieveTranscript: # Retrying won't solve pass # Mark as empty transcript except json.decoder.JSONDecodeError: logger.warning(f'JSONDecodeError for {video_id}') if os.path.exists(transcript_path): os.remove(transcript_path) # Remove file and try again return get_words(video_id, process, transcript_type, fallback, granularity) # Tried to process it, but it was empty... if download or (process and not os.path.exists(transcript_path)): with open(transcript_path, 'w') as fp: json.dump(raw_transcript_json, fp) if not raw_transcript_json and fallback is not None: return get_words(video_id, process, transcript_type=fallback, fallback=None, granularity=granularity) if raw_transcript_json: processed_transcript = parse_transcript_json( raw_transcript_json, granularity) if filter_words_to_remove: processed_transcript = list( filter(lambda x: x['text'] not in WORDS_TO_REMOVE, processed_transcript)) else: processed_transcript = raw_transcript_json # Either None or [] return processed_transcript # TODO make min_sponsor_segment_length param # TODO rename to extract_segments def extract_sponsors(words, min_sponsor_segment_length=3): if not words: return [] paragraphs = [] current = [] prev_category = None for i in range(len(words) + 1): unimportant = i == len(words) or words[i].get('category') is None if unimportant or words[i].get('category') != prev_category: if current: # Save the current batch paragraphs.append({ 'words': current, 'category': current[-1].get('category'), }) current = [] if not unimportant: # Some useful information to save current.append(words[i]) prev_category = words[i].get('category') # Remove all too short: return list(filter(lambda x: len(x['words']) >= min_sponsor_segment_length, paragraphs)) def clean_text(text): # Replace impossibly long words with a special token # Usually the result of incorrect labelling text = re.sub(r'\w{64,}', CustomTokens.LONG_WORD.value, text) SHORT_HYPHENATED_REGEX = r'\w{1,2}(?:-\w{1,2}){3,}(?:-?\w*)' # Replace hyphenated URLs with special token # For some reason, youtube sometimes transcribes urls in this form: # 'b-a-b-b-e-l-dot-com', 'g-e-t-r-o-m-a-n-com' # not 'e-commerce' text = re.sub(f'{SHORT_HYPHENATED_REGEX}(?:com|org|net)', CustomTokens.HYPHENATED_URL.value, text) # Replace short+hyphenated text with a special token. Of the form: # 'i-i-i-i-i-i-i-i-i-i-i-i', 'b-u-m-f-u-z-z-l-e', 'v-e-r-i-t-a-s-i-u-m', 'do-do-do-do-do' text = re.sub(SHORT_HYPHENATED_REGEX, CustomTokens.SHORT_HYPHENATED.value, text) # Replace URLs with URL_TOKEN URL_REGEX = r'(?:(?:http|https)\:\/\/)?[a-zA-Z0-9\.\/\?\:@\-_=#]+\.(?:[a-zA-Z]){2,6}(?:[a-zA-Z0-9\.\&\/\?\:@\-_=#%])*' text = re.sub(URL_REGEX, CustomTokens.URL.value, text) NUM_REGEX = r'(?:\d+,)*(?:\d*[.])?\d+' # Encode specific numeric words # Of the form: 12%, 12.34% # Usually included in sponsorships text = re.sub(f'{NUM_REGEX}%', CustomTokens.NUMBER_PERCENTAGE.value, text) # Normal numbers, should not have an effect on sponsorship text = re.sub(NUM_REGEX, CustomTokens.NUMBER.value, text) # Replace profanity with special token text = text.replace(PROFANITY_RAW, CustomTokens.PROFANITY.value) text = text.replace(PROFANITY_CONVERTED, CustomTokens.PROFANITY.value) return text.strip() def remove_duplicate_segments(segments): # Algorithm based on SponsorBlock algorithm # https://blog.ajay.app/voting-and-pseudo-randomness-or-sponsorblock-or-youtube-sponsorship-segment-blocker # Find sponsors that are overlapping best = [] for i in segments: similar_segments = [] for j in segments: if jaccard(i['start'], i['end'], j['start'], j['end']) > 0.1: # Some overlap similar_segments.append(j) if similar_segments: best_similar_seg = max(similar_segments, key=lambda item: ( item['locked'], item['votes'], item['views'], item['reputation'] )) if best_similar_seg not in best: best.append(best_similar_seg) if len(segments) != len(best): # Saw some reduction... try again return remove_duplicate_segments(best) return best @dataclass class PreprocessArguments: """ Arguments pertaining to what data we are going to preprocess. """ update_database: bool = field( default=False, metadata={'help': 'Download the raw database.'} ) do_create: bool = field( default=False, metadata={'help': 'Merge sponsor segments into single file'} ) min_votes: int = field( default=0, metadata={'help': 'Minimum number of votes'}) # Downvotes will make this negative. # 1 = At least one positive vote max_segment_duration: float = field( default=180, # 3 minutes # >180 => 2.8% # >200 => 2.1% # >250 => 1.1% # >300 => 0.06% metadata={'help': 'Ignore all segments whose duration in seconds is longer than this value (negative means no limit)'}) min_views: int = field( default=5, metadata={'help': 'Minimum number of views a segment must have to be considered. 0 = show all'}) # min_reputation: int = field( # default=0, metadata={'help': 'Minimum reputation a user must have for the segment to be included'}) min_date: str = field( # default='08/06/2020', # release of v2.0 (https://github.com/ajayyy/SponsorBlock/releases/tag/2.0) # release of v3.0 (https://github.com/ajayyy/SponsorBlock/releases/tag/3.0) default='20/08/2021', # default='01/10/2020', # No more autovote metadata={'help': 'Only use submissions from after this date (inclusive)'}) max_date: str = field( # default='01/01/9999', # Include all default='15/04/2022', metadata={'help': 'Only use videos that have some segment from before this date (exclusive). This allows for videos to have segments be corrected, but ignores new videos (posted after this date) to enter the pool.'}) # max_unseen_date: str = field( # TODO # default='02/03/2022', # metadata={'help': 'Generate test and validation data from `max_date` to `max_unseen_date`'}) # Specify min/max video id for splitting (seen vs. unseen) keep_duplicate_segments: bool = field( default=False, metadata={'help': 'Keep duplicate segments'} ) do_process_database: bool = field( default=False, metadata={'help': 'Process the raw database'} ) do_transcribe: bool = field( default=False, metadata={'help': 'Get transcripts for videos'} ) num_jobs: int = field( default=4, metadata={'help': 'Number of transcripts to download in parallel'}) # overwrite: bool = field( # default=False, metadata={'help': 'Overwrite training, testing and validation data, if present.'} # ) do_generate: bool = field( default=False, metadata={'help': 'Generate labelled data.'} ) do_split: bool = field( default=False, metadata={'help': 'Generate training, testing and validation data.'} ) positive_file: Optional[str] = field( default='sponsor_segments.json', metadata={'help': 'File to output sponsored segments to (a jsonlines file).'} ) negative_file: Optional[str] = field( default='normal_segments.json', metadata={'help': 'File to output normal segments to (a jsonlines file).'} ) percentage_positive: float = field( default=0.5, metadata={'help': 'Ratio of positive (sponsor) segments to include in final output'}) train_split: float = field( default=0.9, metadata={'help': 'Ratio of training data. Value between 0 and 1.'}) # TODO play around with ratios? lower test/validation split? test_split: float = field( default=0.05, metadata={'help': 'Ratio of testing data. Value between 0 and 1.'}) valid_split: float = field( default=0.05, metadata={'help': 'Ratio of validation data. Value between 0 and 1.'}) start_index: int = field(default=None, metadata={ 'help': 'Video to start at.'}) max_videos: int = field(default=None, metadata={ 'help': 'Maximum number of videos to preprocess.'}) max_segments: int = field(default=None, metadata={ 'help': 'Maximum number of segments to produce to preprocess.'}) raw_data_dir: Optional[str] = field( default='raw', metadata={ 'help': 'Raw data directory' }, ) raw_data_file: Optional[str] = field( default='sponsorTimes.csv', metadata={ 'help': 'Raw data file' }, ) min_wps: float = field( default=1.5, metadata={'help': 'Ignore videos with not enough words spoken per second. This is usually indicitive of video whose captions aren\'t English.'}) # 0.1 ~ 1% # 0.4 ~ 2.5% # 0.9 ~ 5% # Mirrors for database MIRRORS = [ 'https://sponsor.ajay.app/database/sponsorTimes.csv', # Latest 'https://sb-mirror.mchang.xyz/sponsorTimes.csv', # 5 minute delay 'https://sb.ltn.fi/database/sponsorTimes.csv', # 5 minute delay ] # TODO only download latest updates/changes def download_file(url, filename): """ Helper method handling downloading large files from `url` to `filename`. Adapted from https://stackoverflow.com/a/42071418 """ chunk_size = 1024 r = requests.get(url, stream=True) total_bytes = int(r.headers['Content-Length']) with open(filename, 'wb') as f, tqdm(unit='B', total=total_bytes) as progress: for chunk in r.iter_content(chunk_size=chunk_size): if chunk: # filter out keep-alive new chunks progress.update(len(chunk)) f.write(chunk) return total_bytes == os.path.getsize(filename) def main(): # Responsible for getting transcrips using youtube_transcript_api, # then labelling it according to SponsorBlock's API logger.setLevel(logging.DEBUG) # Generate final.json from sponsorTimes.csv hf_parser = HfArgumentParser(( PreprocessArguments, DatasetArguments, segment.SegmentationArguments, model_module.ModelArguments, GeneralArguments )) preprocess_args, dataset_args, segmentation_args, model_args, general_args = hf_parser.parse_args_into_dataclasses() raw_dataset_path = os.path.join( preprocess_args.raw_data_dir, preprocess_args.raw_data_file) if preprocess_args.update_database: logger.info('Updating database') for mirror in MIRRORS: logger.info(f'Downloading from {mirror}') if download_file(mirror, raw_dataset_path): break logger.warning('Failed, trying next') os.makedirs(dataset_args.data_dir, exist_ok=True) processed_db_path = os.path.join( dataset_args.data_dir, dataset_args.processed_database) # TODO process all valid possible items and then do filtering only later @lru_cache(maxsize=1) def read_db(): # if not preprocess_args.overwrite and os.path.exists(processed_db_path): # logger.info( # 'Using cached processed database (use `--overwrite` to avoid this behaviour).') # with open(processed_db_path) as fp: # return json.load(fp) logger.info('Processing raw database') db = {} allowed_categories = list(map(str.lower, CATGEGORY_OPTIONS)) with open(raw_dataset_path, newline='') as csvfile: reader = csv.DictReader(csvfile) for line in reader: # Never show: if line['service'] != 'YouTube': continue if len(line['videoID']) != 11: continue # Invalid youtube video ID if line['category'] not in allowed_categories: continue if line['actionType'] not in ACTION_OPTIONS: continue # Ignore hidden items if line['hidden'] == '1' or line['shadowHidden'] == '1': continue # Skip those that aren't highly voted votes = int(line['votes']) if votes < preprocess_args.min_votes: continue locked = line['locked'] == '1' reputation = float(line['reputation']) # if reputation < preprocess_args.min_reputation: # continue # TODO add back? # Problems like mGVn1wCkBrE # TODO ignore if over max_duration if line['videoID'] not in db: db[line['videoID']] = [] db[line['videoID']].append({ 'uuid': line['UUID'], 'start': float(line['startTime']), 'end': float(line['endTime']), 'votes': votes, 'locked': locked, 'views': int(line['views']), 'submission_time': float(line['timeSubmitted'])/1e3, 'reputation': reputation, 'category': line['category'], 'action': line['actionType'], }) # First, remove videos that contain a full-video label # (may confuse model since disclaimers and such aren't labelled) # Must do it here before removing duplicate segments for key in list(db): if any(x['action'] == 'full' for x in db[key]): del db[key] # Remove duplicate sponsor segments by choosing best (most votes) if not preprocess_args.keep_duplicate_segments: logger.info('Remove duplicate segments') for key in db: db[key] = remove_duplicate_segments(db[key]) # We now remove whole videos from the list # Helps with obtaining "fully-labelled" videos min_date = datetime.strptime(preprocess_args.min_date, '%d/%m/%Y') max_date = datetime.strptime(preprocess_args.max_date, '%d/%m/%Y') for key in list(db): if preprocess_args.max_segment_duration >= 0 and any(x['end'] - x['start'] > preprocess_args.max_segment_duration for x in db[key]): # Remove videos that have at least one segment that is longer than # the maximum allowed segment duration. This avoids introducing # segments into training that might contain ignored context (since # they are too long, so the middle might be normal content) del db[key] elif any(datetime.fromtimestamp(x['submission_time']) < min_date for x in db[key]): # Remove videos where any of its segments were submitted before min_date # (essentially removes videos uploaded before min_date) # Prevents issues where some segments of a video are excluded del db[key] elif all(datetime.fromtimestamp(x['submission_time']) > max_date for x in db[key]): # Remove videos where all of its segments were submitted after max_date # (essentially removes videos uploaded after max_date) # Allows for segments to be corrected for past videos del db[key] elif any(not x['locked'] and x['views'] < preprocess_args.min_views for x in db[key]): # Remove videos where any of its non-locked segments do not have enough views # (essentially skips videos that have not been fully watched/reviewed) # Always include segments locked by VIPs, regardless of view count del db[key] logger.info(f'Saved {len(db)} videos') with open(processed_db_path, 'w') as fp: json.dump(db, fp) return db if preprocess_args.do_process_database: read_db() # 'videoID', 'startTime', 'endTime', 'votes', 'locked', 'incorrectVotes', 'UUID', # 'userID', 'timeSubmitted', 'views', 'category', 'actionType', 'service', 'videoDuration', # 'hidden', 'reputation', 'shadowHidden', 'hashedVideoID', 'userAgent', 'description' if preprocess_args.do_transcribe: logger.info('Collecting videos') parsed_database = read_db() # Remove transcripts already processed finished = set(x.split('.')[0] for x in os.listdir( 'transcripts/auto/') + os.listdir('transcripts/manual/')) video_ids = list(parsed_database.keys() - finished) # https://stackoverflow.com/a/63495323 import concurrent POLL_INTERVAL = 0.1 # Wrap get words function to return video_id after completion def get_words_wrapper(video_id): get_words(video_id) return video_id logger.info('Setting up ThreadPoolExecutor') with concurrent.futures.ThreadPoolExecutor(max_workers=preprocess_args.num_jobs) as pool, \ tqdm(total=len(video_ids)) as progress: all_futures = (pool.submit(get_words_wrapper, video_id) for video_id in video_ids) to_process = set(itertools.islice( all_futures, preprocess_args.num_jobs)) try: while to_process: just_finished, to_process = concurrent.futures.wait( to_process, timeout=POLL_INTERVAL) to_process |= set(itertools.islice( all_futures, len(just_finished))) for d in just_finished: progress.set_description(f'Processed {d.result()}') progress.update() except KeyboardInterrupt: logger.info( 'Gracefully shutting down: Cancelling unscheduled tasks') # only futures that are not done will prevent exiting for future in to_process: future.cancel() logger.info('Waiting for in-progress tasks to complete') concurrent.futures.wait(to_process, timeout=None) logger.info('Cancellation successful') final_path = os.path.join( dataset_args.data_dir, dataset_args.processed_file) if preprocess_args.do_create: logger.info('Create final data') final_data = {} parsed_database = read_db() transcribed = set(x.split('.')[0] for x in os.listdir( 'transcripts/auto/') + os.listdir('transcripts/manual/')) # Only consider videos that have been transcribed already video_ids = parsed_database.keys() & transcribed with tqdm(total=len(video_ids)) as progress: for index, video_id in enumerate(video_ids): if preprocess_args.max_videos is not None and index >= preprocess_args.max_videos: break progress.set_description(f'Processing {video_id}') progress.update() video_words = get_words(video_id, process=False) if not video_words: continue final_vid_segs = [] # Only add segments with high enough wps for seg in parsed_database[video_id]: segment_words = segment.extract_segment( video_words, seg['start'], seg['end']) if len(segment_words) <= 1: continue # Useless to add segment since no words # duration = segment.word_end(segment_words[-1]) - segment.word_start(segment_words[0]) duration = seg['end'] - seg['start'] wps = len(segment_words)/duration if duration > 0 else 0 # print(video_id, wps) if wps < preprocess_args.min_wps: # Skip sponsor segments without many words # e.g. music ads with some words on each side # progress.set_description(f'Skipping bad segment in {video_id} (wps={wps})') continue final_vid_segs.append(seg) if final_vid_segs: final_data[video_id] = final_vid_segs # Save data with open(final_path, 'w') as fp: json.dump(final_data, fp) # final_data = preprocess( # raw_dataset_path, final_path, preprocess_args.min_votes) # # TODO save metadata in final.json? elif os.path.exists(final_path): # Already exists logging.info(f'{final_path} exists, opening file') with open(final_path) as fp: final_data = json.load(fp) logging.info(f'Found {len(final_data)} videos') else: return # Do not continue # TODO shuffle final_data # if not os.path.exists(excess_path) or preprocess_args.overwrite # TODO use overwrite param positive_file = os.path.join( dataset_args.data_dir, preprocess_args.positive_file) negative_file = os.path.join( dataset_args.data_dir, preprocess_args.negative_file) if preprocess_args.do_generate: logger.info('Generating') # max_videos=preprocess_args.max_videos, # max_segments=preprocess_args.max_segments, # , max_videos, max_segments from model import get_model_tokenizer model, tokenizer = get_model_tokenizer(model_args, general_args) # TODO # count_videos = 0 # count_segments = 0 data = final_data.items() start_index = preprocess_args.start_index or 0 end_index = (preprocess_args.max_videos or len(data)) + start_index data = list(itertools.islice(data, start_index, end_index)) write_mode = 'w' # if preprocess_args.overwrite else 'a' with open(positive_file, write_mode, encoding='utf-8') as positive, \ open(negative_file, write_mode, encoding='utf-8') as negative, \ tqdm(data) as progress: for offset, (video_id, sponsor_segments) in enumerate(data): progress.set_description(f'Processing {video_id}') progress.update() # Use chunk granularity to improve manual transcripts words = get_words(video_id, process=False, granularity='chunk') if not words: continue if len(words) <= 1: continue segments = segment.generate_labelled_segments( words, tokenizer, segmentation_args, sponsor_segments) if not segments: continue for seg in segments: seg_start = segment.word_start(seg[0]) seg_end = segment.word_end(seg[-1]) duration = seg_end - seg_start wps = len(seg)/duration if duration > 0 else 0 # Ignore segments with "not enough words" in the transcript # Must do here since this includes non-sponsor segments if wps < preprocess_args.min_wps: continue d = { # 'video_index': offset + start_index, 'video_id': video_id, # 'uuid': video_id, # TODO add uuid 'text': ' '.join(x['cleaned'] for x in seg), 'start': seg_start, 'end': seg_end, } extracted_segments = extract_sponsors(seg) if extracted_segments: extracted_texts = [] for s in extracted_segments: w = ' '.join(q['cleaned'] for q in s['words']) category = s['category'].upper() extracted_texts.append( f'{START_SEGMENT_TEMPLATE.format(category)} {w} {END_SEGMENT_TEMPLATE.format(category)}' ) d['extracted'] = f' {CustomTokens.BETWEEN_SEGMENTS.value} '.join( extracted_texts) print(json.dumps(d), file=positive) else: d['extracted'] = CustomTokens.NO_SEGMENT.value print(json.dumps(d), file=negative) if preprocess_args.do_split: logger.info('Splitting') logger.info('Read files') with open(positive_file, encoding='utf-8') as positive: sponsors = positive.readlines() with open(negative_file, encoding='utf-8') as negative: non_sponsors = negative.readlines() logger.info('Shuffle') random.shuffle(sponsors) random.shuffle(non_sponsors) logger.info('Calculate ratios') # Ensure correct ratio of positive to negative segments percentage_negative = 1 - preprocess_args.percentage_positive if preprocess_args.percentage_positive * len(sponsors) > len(non_sponsors): # Negative is limiting z = int(preprocess_args.percentage_positive / percentage_negative * len(non_sponsors)) # excess = sponsors[z:] sponsors = sponsors[:z] else: # Positive is limiting z = int(percentage_negative / preprocess_args.percentage_positive * len(sponsors)) # excess = non_sponsors[z:] non_sponsors = non_sponsors[:z] logger.info('Join') all_labelled_segments = sponsors + non_sponsors random.shuffle(all_labelled_segments) # TODO split based on video ids logger.info('Split') ratios = [preprocess_args.train_split, preprocess_args.test_split, preprocess_args.valid_split] train_data, test_data, valid_data = split( all_labelled_segments, ratios) splits = { dataset_args.train_file: train_data, dataset_args.test_file: test_data, dataset_args.validation_file: valid_data } # Output training, testing and validation data for name, items in splits.items(): outfile = os.path.join(dataset_args.data_dir, name) with open(outfile, 'w', encoding='utf-8') as fp: fp.writelines(items) classifier_splits = { dataset_args.c_train_file: train_data, dataset_args.c_test_file: test_data, dataset_args.c_validation_file: valid_data } none_category = CATEGORIES.index(None) # Output training, testing and validation data for name, items in classifier_splits.items(): outfile = os.path.join(dataset_args.data_dir, name) with open(outfile, 'w', encoding='utf-8') as fp: for item in items: parsed_item = json.loads(item) # TODO add uuid matches = extract_sponsor_matches_from_text( parsed_item['extracted']) if matches: for match in matches: print(json.dumps({ 'text': match['text'], 'label': CATEGORIES.index(match['category']) }), file=fp) else: print(json.dumps({ 'text': parsed_item['text'], 'label': none_category }), file=fp) logger.info('Write') # Save excess items # excess_path = os.path.join( # dataset_args.data_dir, dataset_args.excess_file) # if not os.path.exists(excess_path) or preprocess_args.overwrite: # with open(excess_path, 'w', encoding='utf-8') as fp: # fp.writelines(excess) # else: # logger.info(f'Skipping {dataset_args.excess_file}') logger.info( f'Finished splitting: {len(sponsors)} sponsors, {len(non_sponsors)} non sponsors') def split(arr, ratios): """Split array according to ratios. Sum of ratios should be <= 1""" to_return = [] cumulative_sum = 0 for r in ratios: current = cumulative_sum cumulative_sum += r * len(arr) to_return.append(arr[int(current):int(cumulative_sum)]) return to_return if __name__ == '__main__': main()