bert_api.py

# bert-api/bert_api.py
# Origen: D:\XDev.Projects\MEHEARSAL-NLP-CONTROLLER-RMIT\mehearsal_v5_and_v6\bert-api\bert_api.py
# Adaptaciones mínimas para HF Spaces Docker: import os + MODEL_NAME desde env + puerto desde env + debug=False.
# Lógica de extract_entities() y upgrade de intents intacta respecto al fuente.
import os
from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from pathlib import Path
import re


# Get the project root (parent of bert-api folder)
PROJECT_ROOT = Path(__file__).resolve().parent.parent

app = Flask(__name__)
CORS(app)  # Allow requests from Next.js

# Load model and tokenizer from Hugging Face.
# MODEL_NAME es configurable por env var (default: v6). HF_TOKEN se lee automáticamente
# por huggingface_hub para autenticar el acceso al modelo privado.

MODEL_NAME = os.environ.get('BERT_MODEL_NAME', 'MuseSceneLab/mehearsal-nlp-v6')

print(f"Loading model from Hugging Face: {MODEL_NAME}...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

model.eval()

print("Model loaded successfully!")


# Intent to target/params mapping (simplified - you can expand this)
INTENT_METADATA = {
    'MUTE_INSTRUMENT': {'requires_target': True, 'target_type': 'instrument'},
    'NOT_A_COMMAND': {'requires_target': False},
    'UNMUTE_INSTRUMENT': {'requires_target': True, 'target_type': 'instrument'},
    'SOLO_INSTRUMENT': {'requires_target': True, 'target_type': 'instrument'},
    'MUTE_ALL': {'requires_target': False},
    'VOLUME_ADJUST_RELATIVE': {'requires_target': True, 'target_type': 'instrument'},
    'VOLUME_SET_ABSOLUTE': {'requires_target': True, 'target_type': 'instrument'},
    'TEMPO_SET_ABSOLUTE': {'requires_target': False},
    'TEMPO_ADJUST_RELATIVE': {'requires_target': False},
    'TEMPO_GRADUAL': {'requires_target': False},
    'TEMPO_STYLE_MARKING': {'requires_target': False},
    'JUMP_TO_BAR': {'requires_target': False},
    'JUMP_TO_SECTION': {'requires_target': True, 'target_type': 'section'},
    'LOOP_BARS': {'requires_target': False},
    'LOOP_SECTION': {'requires_target': True, 'target_type': 'section'},
    'STOP_LOOP': {'requires_target': False},
}

def extract_entities(text, intent):
    """
    Simple entity extraction based on keywords.
    This is a basic implementation - you might want to enhance this.
    """

    text_lower = text.lower()

    # Instrument keywords
    instruments = ['drums', 'guitar', 'bass', 'piano', 'synth', 'vocals','strings',
                   'batería', 'guitarra', 'bajo']

    # Section keywords
    sections = ['verse', 'chorus', 'bridge', 'outro',
                'introducción', 'verso', 'estribillo', 'puente', 'intro']

    # Detect language
    locale = 'es' if any(word in text_lower for word in ['batería', 'guitarra', 'bajo']) else 'en'

    # Extract target
    target = None
    metadata = INTENT_METADATA.get(intent, {})

    if metadata.get('requires_target'):
        if metadata.get('target_type') == 'instrument':
            for inst in instruments:
                if inst in text_lower:
                    target = inst
                    break
        elif metadata.get('target_type') == 'section':
            # Use word boundary regex to avoid partial matches
            for sec in sections:
                # Match whole word only - use word boundaries
                if re.search(r'\b' + re.escape(sec) + r'\b', text_lower):
                    target = sec
                    break
                    # If no section found, check for instruments (for cases like "loop the guitar")
            if not target:
                for inst in instruments:
                    if re.search(r'\b' + re.escape(inst) + r'\b', text_lower):
                        target = inst
                        break

    # Extract parameters (basic implementation)
    params_json = {}

    # Extract numbers for tempo, volume, bars, etc.
    numbers = re.findall(r'\d+(?:\.\d+)?|\.\d+', text)

    if 'TEMPO_SET_ABSOLUTE' in intent and numbers:
        params_json['bpm'] = int(numbers[0])
    elif 'TEMPO_ADJUST_RELATIVE' in intent:
        # Determine direction
        if 'up' in text_lower or 'increase' in text_lower or 'faster' in text_lower or 'sube' in text_lower or 'más' in text_lower:
            params_json['direction'] = 'up'
        elif 'down' in text_lower or 'decrease' in text_lower or 'slower'  in text_lower or 'off'  in text_lower or 'ease' in text_lower or  'baja' in text_lower or 'menos' in text_lower:
            params_json['direction'] = 'down'

        # Extract percentage if provided
        if numbers:
            params_json['bpm_change_percent'] = int(numbers[0])
        else:
            # Default percentage if not specified
            params_json['bpm_change_percent'] = 10
    elif 'TEMPO_FACTOR' in intent:
        # Handle tempo multipliers
        if 'double' in text_lower or 'twice' in text_lower or 'doble' in text_lower:
            params_json['multiplier'] = 2.0
        elif 'triple' in text_lower or 'three times' in text_lower:
            params_json['multiplier'] = 3.0
        elif 'half' in text_lower or 'mitad' in text_lower:
            params_json['multiplier'] = 0.5
        elif 'quarter' in text_lower or 'cuarto' in text_lower:
            params_json['multiplier'] = 0.25
        elif numbers:
            # If a number is found, use it as multiplier (e.g., "4 times faster")
            multiplier = float(numbers[0])
            # Check if it's a fraction (e.g., "1/2 speed")
            if '/' in text:
                fraction_match = re.search(r'(\d+)/(\d+)', text)
                if fraction_match:
                    numerator = float(fraction_match.group(1))
                    denominator = float(fraction_match.group(2))
                    multiplier = numerator / denominator
            params_json['multiplier'] = multiplier
        else:
            # Default multiplier
            params_json['multiplier'] = 1.0
    elif 'TEMPO_STYLE_MARKING' in intent:
        # Common tempo markings (Italian musical terms)
        tempo_markings = {
            'grave': 'Grave',
            'largo': 'Largo',
            'lento': 'Lento',
            'adagio': 'Adagio',
            'andante': 'Andante',
            'moderato': 'Moderato',
            'allegretto': 'Allegretto',
            'allegro': 'Allegro',
            'vivace': 'Vivace',
            'presto': 'Presto',
            'prestissimo': 'Prestissimo'
        }

        # Find which tempo marking is in the text
        for marking_lower, marking_proper in tempo_markings.items():
            if marking_lower in text_lower:
                params_json['style_marking'] = marking_proper
                break

        # If no specific marking found, default to empty
        if 'style_marking' not in params_json:
            params_json['style_marking'] = ''
    elif 'TEMPO_GRADUAL' in intent:
        # Determine direction (accelerando = speed up, ritardando/rallentando = slow down)
        if any(word in text_lower for word in ['accelerando', 'accel', 'speed up', 'faster', 'acelerar', "crescendo"]):
            params_json['direction'] = 'up'
        elif any(word in text_lower for word in ['ritardando', 'rallentando', 'rit', 'rall', 'slow down', 'slower', 'ralentizar', "diminuendo"]):
            params_json['direction'] = 'down'
        else:
            params_json['direction'] = 'down'  # Default to slowing down

        # Extract bpm_change_percent (default to 10 if not specified)
        percent_match = re.search(r'(\d+\.?\d*)\s*(?:%|percent|por ciento)', text_lower)
        if percent_match:
            params_json['bpm_change_percent'] = float(percent_match.group(1))
        else:
            params_json['bpm_change_percent'] = 10

        # Extract number of bars (look for "over X bars", "in X bars", etc.)
        bars_match = re.search(r'(?:over|in|during|en|durante)\s*(\d+)\s*(?:bar|bars|measure|measures|compás|compases)', text_lower)
        if bars_match:
            params_json['bars'] = int(bars_match.group(1))
        elif numbers:
            # If no explicit "bars" keyword, use first number found
            params_json['bars'] = int(numbers[0])
        else:
            params_json['bars'] = 4  # Default to 4 bars
    elif 'LOOP_BARS' in intent or intent == 'LOOP':
        # Look for bar range patterns like "19-21", "19 to 21", "19 through 21"
        range_match = re.search(r'(?:bar|bars|compás|compases)?\s*(\d+)\s*(?:-|to|through|hasta|a)\s*(\d+)', text_lower)

        if range_match:
            params_json['start_bar'] = int(range_match.group(1))
            params_json['end_bar'] = int(range_match.group(2))
        elif numbers and len(numbers) >= 2:
            # If two numbers found without explicit range pattern
            params_json['start_bar'] = int(numbers[0])
            params_json['end_bar'] = int(numbers[1])
        else:
            # No bar range specified
            params_json['start_bar'] = None
            params_json['end_bar'] = None
    elif 'LOOP_SECTION' in intent:
        # Extract bars if specified
        bars_match = re.search(r'(?:for|during|en|durante)?\s*(\d+)\s*(?:bar|bars|measure|measures|compás|compases)', text_lower)

        if bars_match:
            params_json['bars'] = int(bars_match.group(1))
        else:
            params_json['bars'] = None

        # Set start_bar and end_bar to null for section loops
        params_json['start_bar'] = None
        params_json['end_bar'] = None
    elif 'JUMP_TO_BAR' in intent and numbers:
        params_json['to_bar'] = int(float(numbers[0]))
    elif 'JUMP_RELATIVE' in intent:
        # Determine direction
        if any(word in text_lower for word in ['ahead', 'forward', 'next', 'adelante', 'siguiente']):
            params_json['direction'] = 'up'
        elif any(word in text_lower for word in ['back', 'backward', 'previous', 'atrás', 'anterior']):
            params_json['direction'] = 'down'
        else:
            params_json['direction'] = 'up'  # Default to forward

        # Extract number of bars
        bars_match = re.search(r'(\d+\.?\d*)\s*(?:bar|bars|measure|measures|compás|compases)', text_lower)
        if bars_match:
            params_json['relative_bars'] = int(float(bars_match.group(1)))
        elif numbers:
            params_json['relative_bars'] = int(float(numbers[0]))
        else:
            params_json['relative_bars'] = 1  # Default to 1 bar
    elif 'VOLUME' in intent and numbers:
        params_json['vol'] = int(float(numbers[0]))

        # Detect direction from keywords
        if any(word in text_lower for word in ['up', 'increase', 'boost', 'raise', 'louder', 'sube', 'aumenta', 'más fuerte']):
            params_json['direction'] = 'up'
        elif any(word in text_lower for word in ['down', 'decrease', 'lower', 'reduce', 'quieter', 'baja', 'reduce', 'más bajo']):
            params_json['direction'] = 'down'
        else:
            # If no explicit direction, default based on context or leave it out
            params_json['direction'] = 'up'  # Default assumption
    elif 'MUTE_ALL' in intent:
        # Detect if it's turning mute all on or off
        if any(word in text_lower for word in ['unmute', 'activate', 'turn on', 'on', 'back', 'enable', 'activar', 'encender']):
            params_json['mute_all'] = 'off'
        else:
            # Default to muting (turning on)
            params_json['mute_all'] = 'on'

    return target, locale, params_json

@app.route('/classify', methods=['POST'])
def classify():
    try:
        data = request.json
        utterance = data.get('utterance', '')

        if not utterance:
            return jsonify({'error': 'No utterance provided'}), 400

        # Tokenize input
        inputs = tokenizer(utterance, return_tensors="pt", truncation=True, max_length=512)

        # Run inference
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            predicted_class = torch.argmax(logits, dim=1).item()
            confidence = torch.softmax(logits, dim=1)[0][predicted_class].item()

        # Get intent label
        intent_label = model.config.id2label[predicted_class]

        # Fallback: Upgrade generic intents when specific targets are detected
        text_lower = utterance.lower()

        # Section keywords
        sections = ['verse', 'chorus', 'bridge', 'outro', 'intro',
                    'introducción', 'verso', 'estribillo', 'puente']

        # Instrument keywords
        instruments = ['drums', 'guitar', 'vocals','bass', 'piano', 'synth', 'strings',
                    'batería', 'guitarra', 'bajo']

        # Check for sections
        has_section = any(re.search(r'\b' + re.escape(sec) + r'\b', text_lower) for sec in sections)

        # Check for instruments
        has_instrument = any(re.search(r'\b' + re.escape(inst) + r'\b', text_lower) for inst in instruments)

        # Check for bar/measure navigation keywords
        has_bar_navigation = any(word in text_lower for word in ['go to', 'move to', 'jump to', 'ir a', 'mover a', 'saltar a']) and \
                            any(word in text_lower for word in ['bar', 'measure', 'compás'])

        # Check for volume keywords
        has_volume = any(word in text_lower for word in ['volume', 'volumen', 'loud', 'quiet', 'louder', 'quieter', 'más fuerte', 'más bajo'])

        # Upgrade intents based on detected entities
        if intent_label == 'LOOP' and has_section:
            intent_label = 'LOOP_SECTION'
        elif intent_label == 'MUTE' and has_instrument:
            intent_label = 'MUTE_INSTRUMENT'
        elif intent_label == 'UNMUTE' and has_instrument:
            intent_label = 'UNMUTE_INSTRUMENT'
        elif intent_label == 'SOLO' and has_instrument:
            intent_label = 'SOLO_INSTRUMENT'
        elif intent_label in ['VOLUME_SET_ABSOLUTE', 'VOLUME_ADJUST_RELATIVE'] and has_bar_navigation:
            intent_label = 'JUMP_TO_BAR'
        elif intent_label == 'TEMPO_ADJUST_RELATIVE' and has_volume:
            # If model predicts tempo but volume keyword is present, correct to volume
            if has_instrument:
                intent_label = 'VOLUME_ADJUST_RELATIVE'
            else:
                intent_label = 'VOLUME_ADJUST_RELATIVE'

        # Extract entities and parameters
        target, locale, params_json = extract_entities(utterance, intent_label)

        # Build response
        result = {
            'intent_label': intent_label,
            'locale': locale,
            'target': target,
            'params_json': params_json,
            'confidence': round(confidence * 100, 2),
            'model': 'DistilBERT'
        }

        return jsonify(result)

    except Exception as e:
        print(f"Error: {str(e)}")
        return jsonify({
            'intent_label': 'UNKNOWN',
            'locale': 'en',
            'target': None,
            'params_json': {'error': str(e)},
            'model': 'DistilBERT'
        }), 500

@app.route('/health', methods=['GET'])
def health():
    return jsonify({'status': 'healthy', 'model': MODEL_NAME})

if __name__ == '__main__':
    # Puerto adaptado para HF Spaces (default 7860). Debug off para producción.
    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)), debug=False)