src/predict_server.py

#!/usr/bin/env python3
import os
import sys
import logging
import tempfile
from flask import Flask, request, jsonify
from transformers import HfArgumentParser

# Add src directory to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from predict import predict
from preprocess import parse_srt_file
from segment import SegmentationArguments
from shared import GeneralArguments
from model import get_model_tokenizer_classifier, InferenceArguments

# Configure logging
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Initialize Flask app
app = Flask(__name__)

# Global variables to store model, tokenizer, and classifier
model = None
tokenizer = None
classifier = None
segmentation_args = None
inference_args = None


def parse_srt_content(srt_content):
    """Parse SRT content from string and return words with timestamps"""
    import re
    
    words = []
    
    # Handle different line endings
    srt_content = srt_content.replace('\r\n', '\n').replace('\r', '\n')
    
    # Split by double newline to get individual subtitle blocks
    # Also handle cases where there might be extra whitespace
    blocks = re.split(r'\n\s*\n', srt_content.strip())
    
    logger.debug(f"Found {len(blocks)} subtitle blocks")
    
    for i, block in enumerate(blocks):
        lines = block.strip().split('\n')
        logger.debug(f"Block {i}: {len(lines)} lines")
        
        if len(lines) >= 3:
            # Skip the index line (first line)
            # Parse timestamp line (second line)
            timestamp_line = lines[1].strip()
            logger.debug(f"Timestamp line: {timestamp_line}")
            
            match = re.match(r'(\d{2}):(\d{2}):(\d{2}),(\d{3})\s*-->\s*(\d{2}):(\d{2}):(\d{2}),(\d{3})', timestamp_line)
            if match:
                # Convert to seconds
                start_h, start_m, start_s, start_ms = map(int, match.groups()[:4])
                end_h, end_m, end_s, end_ms = map(int, match.groups()[4:])
                
                start_time = start_h * 3600 + start_m * 60 + start_s + start_ms / 1000
                end_time = end_h * 3600 + end_m * 60 + end_s + end_ms / 1000
                
                # Join all text lines (there might be multiple lines of text)
                text = ' '.join(line.strip() for line in lines[2:] if line.strip())
                logger.debug(f"Text: {text}")
                
                # Split text into words
                text_words = text.split()
                if text_words:
                    # Distribute time equally among words
                    duration = end_time - start_time
                    time_per_word = duration / len(text_words) if len(text_words) > 0 else 0
                    
                    for i, word in enumerate(text_words):
                        word_start = start_time + i * time_per_word
                        word_end = start_time + (i + 1) * time_per_word
                        words.append({
                            'text': word,
                            'start': round(word_start, 3),
                            'end': round(word_end, 3)
                        })
            else:
                logger.debug(f"No match for timestamp: {timestamp_line}")
    
    logger.debug(f"Total words parsed: {len(words)}")
    return words


def initialize_model():
    """Initialize model, tokenizer, and classifier"""
    global model, tokenizer, classifier, segmentation_args, inference_args
    
    logger.info("Initializing model...")
    
    # Parse arguments with defaults
    hf_parser = HfArgumentParser((
        InferenceArguments,
        SegmentationArguments,
        GeneralArguments
    ))

    # Create default arguments
    inference_args = InferenceArguments()
    segmentation_args = SegmentationArguments()
    general_args = GeneralArguments()
    
    # Get model, tokenizer, and classifier
    model, tokenizer, classifier = get_model_tokenizer_classifier(
        inference_args, general_args)
    
    logger.info("Model initialized successfully")


def _extract_srt_and_probability():
    """Extract SRT input and optional min_probability from request payload."""
    srt_content = None
    min_probability = request.args.get('min_probability', type=float)

    if request.content_type == 'text/plain':
        srt_content = request.data.decode('utf-8')
        return srt_content, min_probability

    if request.content_type and 'multipart/form-data' in request.content_type:
        if 'file' in request.files:
            file = request.files['file']
            srt_content = file.read().decode('utf-8')
        elif 'srt' in request.form:
            srt_content = request.form['srt']
        return srt_content, min_probability

    if request.content_type == 'application/json':
        data = request.get_json(silent=True) or {}

        if isinstance(data.get('srt'), str):
            srt_content = data['srt']
        elif isinstance(data.get('inputs'), str):
            # Hugging Face style payload: {"inputs": "<srt text>"}
            srt_content = data['inputs']
        elif isinstance(data.get('inputs'), dict) and isinstance(data['inputs'].get('srt'), str):
            # Alternate HF style payload: {"inputs": {"srt": "..."}}
            srt_content = data['inputs']['srt']

        params = data.get('parameters')
        if min_probability is None and isinstance(params, dict):
            min_probability = params.get('min_probability')

    return srt_content, min_probability


def _run_prediction_request():
    srt_content, min_probability = _extract_srt_and_probability()
    if min_probability is None and inference_args is not None:
        min_probability = inference_args.min_probability

    if not srt_content:
        return jsonify({
            'error': 'No SRT content provided. Send text/plain, multipart/form-data with "file" or "srt", or JSON with "srt"/"inputs".'
        }), 400

    logger.debug(f"Received SRT content (first 200 chars): {srt_content[:200]}")

    words = parse_srt_content(srt_content)
    logger.debug(f"Parsed {len(words)} words from SRT")
    if words:
        logger.debug(f"First word: {words[0]}")
        logger.debug(f"Last word: {words[-1]}")

    if not words:
        return jsonify({'error': 'No words found in SRT content'}), 400

    predictions = predict(
        'srt_input',
        model,
        tokenizer,
        segmentation_args,
        words=words,
        classifier=classifier,
        min_probability=min_probability
    )

    output = {'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)

    return jsonify(output)


@app.route('/predict', methods=['POST'])
def predict_endpoint():
    """Primary prediction endpoint."""
    try:
        return _run_prediction_request()
    except Exception as e:
        logger.error(f"Error processing request: {str(e)}")
        return jsonify({'error': str(e)}), 500


@app.route('/', methods=['GET', 'POST'])
def root_predict_endpoint():
    """HF-compatible root endpoint for custom-image Inference Endpoints."""
    try:
        if request.method == 'GET':
            return jsonify({
                'name': 'sponsorblock-ml',
                'status': 'ok',
                'usage': {
                    'predict': {
                        'method': 'POST',
                        'path': '/',
                        'content_types': [
                            'application/json',
                            'text/plain',
                            'multipart/form-data'
                        ],
                        'examples': [
                            {'inputs': '<srt text>'},
                            {'srt': '<srt text>'},
                            {'inputs': {'srt': '<srt text>'}}
                        ]
                    },
                    'health': {
                        'method': 'GET',
                        'path': '/health'
                    }
                }
            })
        return _run_prediction_request()
    except Exception as e:
        logger.error(f"Error processing request: {str(e)}")
        return jsonify({'error': str(e)}), 500


@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({
        'status': 'healthy',
        'model_loaded': model is not None
    })


# Initialize model when module is imported (for Gunicorn)
initialize_model()

if __name__ == '__main__':
    # Only run Flask dev server if running directly
    # For production, use: gunicorn -b 0.0.0.0:5000 src.predict_server:app
    
    # Get port from environment variable or use default
    port = int(os.environ.get('PORT', 5000))
    
    # Run server
    logger.info(f"Starting development server on port {port}")
    app.run(host='0.0.0.0', port=port, debug=False)