src/predict.py

from transformers import HfArgumentParser
from dataclasses import dataclass, field
import logging
import json
from shared import CustomTokens, extract_sponsor_matches, GeneralArguments, seconds_to_time
from segment import (
    generate_segments,
    extract_segment,
    MIN_SAFETY_TOKENS,
    SAFETY_TOKENS_PERCENTAGE,
    word_start,
    word_end,
    SegmentationArguments
)
import preprocess
from errors import TranscriptError
from model import get_model_tokenizer_classifier, InferenceArguments

logging.basicConfig()
logger = logging.getLogger(__name__)




MATCH_WINDOW = 25       # Increase for accuracy, but takes longer: O(n^3)
MERGE_TIME_WITHIN = 8   # Merge predictions if they are within x seconds

# Any prediction whose start time is <= this will be set to start at 0
START_TIME_ZERO_THRESHOLD = 0.08


def filter_and_add_probabilities(predictions, classifier, min_probability):
    """Use classifier to filter predictions"""
    if not predictions:
        return predictions

    # We update the predicted category from the extractive transformer
    # if the classifier is confident enough it is another category

    texts = [
        preprocess.clean_text(' '.join([x['text'] for x in pred['words']]))
        for pred in predictions
    ]
    classifications = classifier(texts)

    filtered_predictions = []
    for prediction, probabilities in zip(predictions, classifications):
        predicted_probabilities = {
            p['label'].lower(): p['score'] for p in probabilities}

        # Get best category + probability
        classifier_category = max(
            predicted_probabilities, key=predicted_probabilities.get)
        classifier_probability = predicted_probabilities[classifier_category]

        if (prediction['category'] not in predicted_probabilities) \
                or (classifier_category != 'none' and classifier_probability > 0.5):  # TODO make param
            # Unknown category or we are confident enough to overrule,
            # so change category to what was predicted by classifier
            prediction['category'] = classifier_category

        if prediction['category'] == 'none':
            continue  # Ignore if categorised as nothing

        prediction['probability'] = predicted_probabilities[prediction['category']]

        if min_probability is not None and prediction['probability'] < min_probability:
            continue  # Ignore if below threshold

        # TODO add probabilities, but remove None and normalise rest
        prediction['probabilities'] = predicted_probabilities

        # if prediction['probability'] < classifier_args.min_probability:
        #     continue

        filtered_predictions.append(prediction)

    return filtered_predictions

class TranscriptError(Exception):
    pass

def predict(identifier, model, tokenizer, segmentation_args, words, classifier=None, min_probability=None):
    # Process words directly (no need to fetch from YouTube)
    if not words:
        raise TranscriptError('No words provided')

    segments = generate_segments(
        words,
        tokenizer,
        segmentation_args
    )

    predictions = segments_to_predictions(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 is not None:
        predictions = filter_and_add_probabilities(
            predictions, classifier, min_probability)

    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_from_texts(texts, model, tokenizer):
    clean_texts = list(map(preprocess.clean_text, texts))
    return predict_sponsor_from_cleaned_texts(clean_texts, model, tokenizer)


def predict_sponsor_from_cleaned_texts(cleaned_texts, model, tokenizer):
    """Given a body of text, predict the words which are part of the sponsor"""
    model_device = next(model.parameters()).device

    decoded_outputs = []
    # Do individually, to avoid running out of memory for long videos
    for cleaned_words in cleaned_texts:
        text = CustomTokens.EXTRACT_SEGMENTS_PREFIX.value + \
            ' '.join(cleaned_words)
        input_ids = tokenizer(text, return_tensors='pt',
                              truncation=True).input_ids.to(model_device)

        # Optimise output length so that we do not generate unnecessarily long texts
        max_out_len = round(min(
            max(
                len(input_ids[0])/SAFETY_TOKENS_PERCENTAGE,
                len(input_ids[0]) + MIN_SAFETY_TOKENS
            ),
            model.model_dim)
        )

        outputs = model.generate(input_ids, max_length=max_out_len)
        decoded_outputs.append(tokenizer.decode(
            outputs[0], skip_special_tokens=True))

    return decoded_outputs


def segments_to_predictions(segments, model, tokenizer):
    predicted_time_ranges = []

    cleaned_texts = [
        [x['cleaned'] for x in cleaned_segment]
        for cleaned_segment in segments
    ]

    sponsorship_texts = predict_sponsor_from_cleaned_texts(
        cleaned_texts, model, tokenizer)

    matches = extract_sponsor_matches(sponsorship_texts)

    for segment_matches, cleaned_batch, segment in zip(matches, cleaned_texts, segments):

        for match in segment_matches:  # one segment might contain multiple sponsors/ir/selfpromos

            matched_text = match['text'].split()

            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'] if range['start'] > START_TIME_ZERO_THRESHOLD else 0
        end_time = range['end']

        if prev_prediction is not None and \
                (start_time <= prev_prediction['end'] <= end_time or    # Merge overlapping segments
                    (range['category'] == prev_prediction['category']   # Merge disconnected segments if same category and within threshold
                        and start_time - prev_prediction['end'] <= MERGE_TIME_WITHIN)):
            # 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
    logger.setLevel(logging.DEBUG)

    import sys
    
    # Extract positional argument for SRT file
    if len(sys.argv) < 2:
        logger.error('Usage: python predict.py <srt_file> [options]')
        return
    
    srt_file = sys.argv[1]
    
    # Remove the positional argument before parsing other arguments
    sys.argv = [sys.argv[0]] + sys.argv[2:]
    
    hf_parser = HfArgumentParser((
        InferenceArguments,
        SegmentationArguments,
        GeneralArguments
    ))
    predict_args, segmentation_args, general_args = hf_parser.parse_args_into_dataclasses()

    # Process SRT file
    from preprocess import parse_srt_file
    import os
    
    if not os.path.exists(srt_file):
        logger.error(f'SRT file not found: {srt_file}')
        return
        
    logger.info(f'Processing SRT file: {srt_file}')
    
    model, tokenizer, classifier = get_model_tokenizer_classifier(
        predict_args, general_args)
    
    try:
        # Parse SRT file to get words with timestamps
        words = parse_srt_file(srt_file)
        
        if not words:
            logger.error('No words found in SRT file')
            return
        
        # Use filename as identifier
        srt_filename = os.path.basename(srt_file)
        
        # Predict segments from parsed words
        predictions = predict(srt_filename, model, tokenizer, segmentation_args,
                              words=words,
                              classifier=classifier,
                              min_probability=predict_args.min_probability)
        
        if not predictions:
            logger.info(f'No predictions found for {srt_filename}')
            return
        
        # Output as JSON
        output = {
            'file': srt_filename,
            'predictions': []
        }
        
        for prediction in predictions:
            pred_data = {
                'time_start': prediction['start'],
                'time_end': prediction['end'],
                'category': prediction.get('category')
            }
            if 'probability' in prediction:
                pred_data['probability'] = prediction['probability']
            output['predictions'].append(pred_data)
        
        print(json.dumps(output, indent=2))
            
    except Exception as e:
        logger.error(f'Error processing SRT file: {e}')
        return


if __name__ == '__main__':
    main()