#!/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_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': ''}, {'inputs': {'srt': ''}} ] }, '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)