Spaces:
Runtime error
Runtime error
File size: 9,570 Bytes
90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 74c1216 5fbdd3c 9b9ffd0 5fbdd3c 9b9ffd0 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 90d1f68 5fbdd3c 9b9ffd0 5fbdd3c 90d1f68 5fbdd3c 74c1216 5fbdd3c 74c1216 90d1f68 74c1216 5fbdd3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 |
from utils import re_findall
from shared import OutputArguments
from typing import Optional
from segment import (
generate_segments,
extract_segment,
SAFETY_TOKENS,
CustomTokens,
word_start,
word_end,
SegmentationArguments
)
import preprocess
from errors import TranscriptError
from model import get_classifier_vectorizer
from transformers import (
AutoModelForSeq2SeqLM,
AutoTokenizer,
HfArgumentParser
)
from transformers.trainer_utils import get_last_checkpoint
from dataclasses import dataclass, field
from shared import device
import logging
def seconds_to_time(seconds):
fractional = round(seconds % 1, 3)
fractional = '' if fractional == 0 else str(fractional)[1:]
h, remainder = divmod(abs(int(seconds)), 3600)
m, s = divmod(remainder, 60)
return f"{'-' if seconds < 0 else ''}{h:02}:{m:02}:{s:02}{fractional}"
@dataclass
class TrainingOutputArguments:
model_path: str = field(
default=None,
metadata={
'help': 'Path to pretrained model used for prediction'}
)
output_dir: Optional[str] = OutputArguments.__dataclass_fields__[
'output_dir']
def __post_init__(self):
if self.model_path is not None:
return
last_checkpoint = get_last_checkpoint(self.output_dir)
if last_checkpoint is not None:
self.model_path = last_checkpoint
else:
raise Exception(
'Unable to find model, explicitly set `--model_path`')
@dataclass
class PredictArguments(TrainingOutputArguments):
video_id: str = field(
default=None,
metadata={
'help': 'Video to predict sponsorship segments for'}
)
SPONSOR_MATCH_RE = fr'(?<={CustomTokens.START_SEGMENT.value})\s*_(?P<category>\S+)\s*(?P<text>.*?)\s*(?={CustomTokens.END_SEGMENT.value}|$)'
MATCH_WINDOW = 25 # Increase for accuracy, but takes longer: O(n^3)
MERGE_TIME_WITHIN = 8 # Merge predictions if they are within x seconds
@dataclass
class ClassifierArguments:
classifier_dir: Optional[str] = field(
default='classifiers',
metadata={
'help': 'The directory that contains the classifier and vectorizer.'
}
)
classifier_file: Optional[str] = field(
default='classifier.pickle',
metadata={
'help': 'The name of the classifier'
}
)
vectorizer_file: Optional[str] = field(
default='vectorizer.pickle',
metadata={
'help': 'The name of the vectorizer'
}
)
min_probability: float = field(
default=0.5, metadata={'help': 'Remove all predictions whose classification probability is below this threshold.'})
def filter_predictions(predictions, classifier_args): # classifier, vectorizer,
"""Use classifier to filter predictions"""
if not predictions:
return predictions
classifier, vectorizer = get_classifier_vectorizer(classifier_args)
transformed_segments = vectorizer.transform([
preprocess.clean_text(' '.join([x['text'] for x in pred['words']]))
for pred in predictions
])
probabilities = classifier.predict_proba(transformed_segments)
filtered_predictions = []
for prediction, probability in zip(predictions, probabilities):
prediction['probability'] = probability[1]
if prediction['probability'] >= classifier_args.min_probability:
filtered_predictions.append(prediction)
# else:
# print('removing segment', prediction)
return filtered_predictions
def predict(video_id, model, tokenizer, segmentation_args, words=None, classifier_args=None):
# Allow words to be passed in so that we don't have to get the words if we already have them
if words is None:
words = preprocess.get_words(video_id)
if not words:
raise TranscriptError('Unable to retrieve transcript')
segments = generate_segments(
words,
tokenizer,
segmentation_args
)
predictions = segments_to_prediction_times(segments, model, tokenizer)
# Add words back to time_ranges
for prediction in predictions:
# Stores words in the range
prediction['words'] = extract_segment(
words, prediction['start'], prediction['end'])
if classifier_args is not None:
predictions = filter_predictions(predictions, classifier_args)
return predictions
def greedy_match(list, sublist):
# Return index and length of longest matching sublist
best_i = -1
best_j = -1
best_k = 0
for i in range(len(list)): # Start position in main list
for j in range(len(sublist)): # Start position in sublist
for k in range(len(sublist)-j, 0, -1): # Width of sublist window
if k > best_k and list[i:i+k] == sublist[j:j+k]:
best_i, best_j, best_k = i, j, k
break # Since window size decreases
return best_i, best_j, best_k
def predict_sponsor_text(text, model, tokenizer):
"""Given a body of text, predict the words which are part of the sponsor"""
input_ids = tokenizer(
f'{CustomTokens.EXTRACT_SEGMENTS_PREFIX.value} {text}', return_tensors='pt', truncation=True).input_ids.to(device())
# Can't be longer than input length + SAFETY_TOKENS or model input dim
max_out_len = min(len(input_ids[0]) + SAFETY_TOKENS, model.model_dim)
outputs = model.generate(input_ids, max_length=max_out_len)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
def predict_sponsor_matches(text, model, tokenizer):
sponsorship_text = predict_sponsor_text(text, model, tokenizer)
if CustomTokens.NO_SEGMENT.value in sponsorship_text:
return []
return re_findall(SPONSOR_MATCH_RE, sponsorship_text)
def segments_to_prediction_times(segments, model, tokenizer):
predicted_time_ranges = []
# TODO pass to model simultaneously, not in for loop
# use 2d array for input ids
for segment in segments:
cleaned_batch = [preprocess.clean_text(
word['text']) for word in segment]
batch_text = ' '.join(cleaned_batch)
matches = predict_sponsor_matches(batch_text, model, tokenizer)
for match in matches:
matched_text = match['text'].split()
# TODO skip if too short
i1, j1, k1 = greedy_match(
cleaned_batch, matched_text[:MATCH_WINDOW])
i2, j2, k2 = greedy_match(
cleaned_batch, matched_text[-MATCH_WINDOW:])
extracted_words = segment[i1:i2+k2]
if not extracted_words:
continue
predicted_time_ranges.append({
'start': word_start(extracted_words[0]),
'end': word_end(extracted_words[-1]),
'category': match['category']
})
# Necessary to sort matches by start time
predicted_time_ranges.sort(key=word_start)
# Merge overlapping predictions and sponsorships that are close together
# Caused by model having max input size
prev_prediction = None
final_predicted_time_ranges = []
for range in predicted_time_ranges:
start_time = range['start']
end_time = range['end']
if prev_prediction is not None and range['category'] == prev_prediction['category'] and (
start_time <= prev_prediction['end'] <= end_time or start_time -
prev_prediction['end'] <= MERGE_TIME_WITHIN
):
# Ending time of last segment is in this segment or c, so we extend last prediction range
final_predicted_time_ranges[-1]['end'] = end_time
else: # No overlap, is a new prediction
final_predicted_time_ranges.append({
'start': start_time,
'end': end_time,
'category': range['category']
})
prev_prediction = range
return final_predicted_time_ranges
def main():
# Test on unseen data
logging.getLogger().setLevel(logging.DEBUG)
hf_parser = HfArgumentParser((
PredictArguments,
SegmentationArguments,
ClassifierArguments
))
predict_args, segmentation_args, classifier_args = hf_parser.parse_args_into_dataclasses()
if predict_args.video_id is None:
print('No video ID supplied. Use `--video_id`.')
return
model = AutoModelForSeq2SeqLM.from_pretrained(predict_args.model_path)
model.to(device())
tokenizer = AutoTokenizer.from_pretrained(predict_args.model_path)
predict_args.video_id = predict_args.video_id.strip()
predictions = predict(predict_args.video_id, model, tokenizer,
segmentation_args) # TODO add back , classifier_args=classifier_args
video_url = f'https://www.youtube.com/watch?v={predict_args.video_id}'
if not predictions:
print('No predictions found for', video_url)
return
print(len(predictions), 'predictions found for', video_url)
for index, prediction in enumerate(predictions, start=1):
print(f'Prediction #{index}:')
print('Text: "',
' '.join([w['text'] for w in prediction['words']]), '"', sep='')
print('Time:', seconds_to_time(
prediction['start']), '-->', seconds_to_time(prediction['end']))
print('Probability:', prediction.get('probability'))
print('Category:', prediction.get('category'))
print()
if __name__ == '__main__':
main()
|