Update app.py

app.py

@@ -3,11 +3,9 @@ import torch
 import logging
 import gc
 import time
-import signal
 from pathlib import Path
-from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
+from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor, WhisperProcessor, WhisperForConditionalGeneration
 import librosa
-from functools import wraps
 
 # Try to import flash attention, but don't fail if not available
 try:
@@ -22,25 +20,6 @@ except ImportError:
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-def timeout_handler(signum, frame):
-    raise TimeoutError("Operation timed out")
-
-def with_timeout(seconds):
-    def decorator(func):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            # Set the signal handler and a timeout alarm
-            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
-            signal.alarm(seconds)
-            try:
-                result = func(*args, **kwargs)
-            finally:
-                signal.alarm(0)  # Disable the alarm
-                signal.signal(signal.SIGALRM, old_handler)
-            return result
-        return wrapper
-    return decorator
-
 class OptimizedWhisperApp:
     def __init__(self):
         self.pipe = None
@@ -57,35 +36,105 @@ class OptimizedWhisperApp:
             "ilsp/whisper_greek_dialect_of_lesbos"
         ]
     
+    def is_fine_tuned_model(self, model_name):
+        """Check if this is a fine-tuned model that might need special handling"""
+        fine_tuned_indicators = [
+            "ilsp/",
+            "fine",
+            "dialect",
+            "custom",
+        ]
+        return any(indicator in model_name.lower() for indicator in fine_tuned_indicators)
+    
+    def create_pipe_for_fine_tuned(self, model_name):
+        """Special handling for fine-tuned models"""
+        try:
+            logger.info(f"Loading fine-tuned model: {model_name}")
+            
+            # Device selection - be more conservative for fine-tuned models
+            if torch.cuda.is_available():
+                device = "cuda:0"
+                torch_dtype = torch.float32  # Use float32 for stability
+            else:
+                device = "cpu"
+                torch_dtype = torch.float32
+            
+            logger.info(f"Using device: {device}, dtype: {torch_dtype}")
+            
+            # Try to load as Whisper model first
+            try:
+                logger.info("Attempting to load as WhisperForConditionalGeneration...")
+                model = WhisperForConditionalGeneration.from_pretrained(
+                    model_name,
+                    torch_dtype=torch_dtype,
+                    low_cpu_mem_usage=True,
+                    cache_dir="./cache"
+                )
+                processor = WhisperProcessor.from_pretrained(model_name)
+                logger.info("Successfully loaded as Whisper model")
+            except Exception as e:
+                logger.info(f"Whisper loading failed: {e}, trying AutoModel...")
+                model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                    model_name,
+                    torch_dtype=torch_dtype,
+                    low_cpu_mem_usage=True,
+                    use_safetensors=False,  # Fine-tuned models might not have safetensors
+                    cache_dir="./cache"
+                )
+                processor = AutoProcessor.from_pretrained(model_name)
+            
+            model.to(device)
+            logger.info("Model moved to device")
+            
+            # Create pipeline with conservative settings
+            pipe = pipeline(
+                "automatic-speech-recognition",
+                model=model,
+                tokenizer=processor.tokenizer,
+                feature_extractor=processor.feature_extractor,
+                torch_dtype=torch_dtype,
+                device=device,
+                chunk_length_s=30,  # Fixed chunk length for fine-tuned models
+            )
+            
+            logger.info("Fine-tuned model pipeline created successfully!")
+            return pipe
+            
+        except Exception as e:
+            logger.error(f"Failed to create fine-tuned model pipeline: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return None
+    
     def create_pipe(self, model_name, use_flash_attention=True):
-        """Create pipeline with better error handling"""
+        """Create pipeline with special handling for fine-tuned models"""
+        
+        # Use special handling for fine-tuned models
+        if self.is_fine_tuned_model(model_name):
+            return self.create_pipe_for_fine_tuned(model_name)
+        
         try:
-            logger.info(f"Starting to load model: {model_name}")
+            logger.info(f"Loading standard model: {model_name}")
             
             # Device selection
             if torch.cuda.is_available():
                 device = "cuda:0"
                 torch_dtype = torch.float16
-                logger.info("Using CUDA device")
             else:
                 device = "cpu"
                 torch_dtype = torch.float32
-                logger.info("Using CPU device")
             
-            # Simpler attention implementation selection
-            attn_implementation = "eager"  # Start with most compatible
+            # Attention implementation - disable for fine-tuned models
+            attn_implementation = "eager"
             if use_flash_attention and FLASH_ATTN_AVAILABLE and is_flash_attn_2_available() and torch.cuda.is_available():
                 try:
                     attn_implementation = "flash_attention_2"
-                    logger.info("Attempting Flash Attention 2")
+                    logger.info("Using Flash Attention 2")
                 except:
                     attn_implementation = "eager"
                     logger.info("Flash Attention 2 failed, using eager")
             
-            logger.info(f"Using attention implementation: {attn_implementation}")
-            
-            # Load model with timeout protection
-            logger.info("Loading model...")
+            # Load model
             model = AutoModelForSpeechSeq2Seq.from_pretrained(
                 model_name,
                 torch_dtype=torch_dtype,
@@ -94,15 +143,12 @@ class OptimizedWhisperApp:
                 attn_implementation=attn_implementation,
                 cache_dir="./cache"
             )
-            logger.info("Model loaded, moving to device...")
             model.to(device)
             
             # Load processor
-            logger.info("Loading processor...")
             processor = AutoProcessor.from_pretrained(model_name)
             
             # Create pipeline
-            logger.info("Creating pipeline...")
             pipe = pipeline(
                 "automatic-speech-recognition",
                 model=model,
@@ -112,11 +158,11 @@ class OptimizedWhisperApp:
                 device=device,
             )
             
-            logger.info("Pipeline created successfully!")
+            logger.info("Standard model pipeline created successfully!")
             return pipe
             
         except Exception as e:
-            logger.error(f"Failed to create pipeline: {e}")
+            logger.error(f"Failed to create standard model pipeline: {e}")
             import traceback
             logger.error(traceback.format_exc())
             return None
@@ -135,7 +181,11 @@ class OptimizedWhisperApp:
                 gc.collect()
             
             try:
-                # Create new pipeline with timeout
+                # Disable flash attention for fine-tuned models
+                if self.is_fine_tuned_model(model_name):
+                    use_flash_attention = False
+                    logger.info("Disabled flash attention for fine-tuned model")
+                
                 self.pipe = self.create_pipe(model_name, use_flash_attention)
                 self.current_model = model_name if self.pipe else None
                 
@@ -153,11 +203,35 @@ class OptimizedWhisperApp:
             logger.info("Model already loaded")
             return True
     
+    def transcribe_audio_fine_tuned(self, audio_file, chunk_length_s=30, batch_size=1):
+        """Special transcription method for fine-tuned models with conservative settings"""
+        try:
+            logger.info("Using fine-tuned model transcription method")
+            
+            # Use very conservative settings for fine-tuned models
+            outputs = self.pipe(
+                audio_file,
+                chunk_length_s=min(chunk_length_s, 30),  # Max 30 seconds
+                batch_size=min(batch_size, 2),  # Max batch size 2
+                return_timestamps=True,
+                generate_kwargs={
+                    "task": "transcribe",
+                    "do_sample": False,  # Deterministic output
+                    "num_beams": 1,      # No beam search
+                    "max_length": 448,   # Conservative max length
+                }
+            )
+            return outputs
+            
+        except Exception as e:
+            logger.error(f"Fine-tuned transcription failed: {e}")
+            raise e
+    
     def transcribe_audio(self, audio_file, model_name="openai/whisper-medium", 
                         language="Automatic Detection", task="transcribe",
                         chunk_length_s=30, batch_size=16, use_flash_attention=True,
                         return_timestamps=True):
-        """Transcribe with comprehensive error handling"""
+        """Transcribe with special handling for fine-tuned models"""
         
         if audio_file is None:
             return "Please upload an audio file", "", ""
@@ -167,40 +241,47 @@ class OptimizedWhisperApp:
             start_time = time.time()
             
             # Load model
-            logger.info(f"Loading model: {model_name}")
             success = self.load_model(model_name, use_flash_attention)
             if not success:
-                return "Failed to load model - check logs for details", "", ""
-            
-            logger.info(f"Processing audio file: {audio_file}")
-            logger.info(f"Settings - Model: {model_name}, Language: {language}, Task: {task}")
-            logger.info(f"Chunk length: {chunk_length_s}s, Batch size: {batch_size}")
-            
-            # Prepare generation kwargs
-            generate_kwargs = {}
-            
-            # Language handling
-            if language != "Automatic Detection" and not model_name.endswith(".en"):
-                language_map = {
-                    "Greek": "greek",
-                    "English": "english", 
-                    "Spanish": "spanish",
-                    "French": "french",
-                    "German": "german",
-                    "Italian": "italian"
-                }
-                lang_code = language_map.get(language, language.lower())
-                generate_kwargs["language"] = lang_code
-                logger.info(f"Set language: {lang_code}")
+                return "Failed to load model", "", ""
             
-            # Task handling
-            if not model_name.endswith(".en"):
-                generate_kwargs["task"] = task
-                logger.info(f"Set task: {task}")
+            logger.info(f"Processing: {audio_file}")
+            logger.info(f"Settings: {model_name}, {language}, {task}")
             
-            # Transcribe with error handling
-            logger.info("Starting transcription process...")
-            try:
+            # Check if this is a fine-tuned model
+            is_fine_tuned = self.is_fine_tuned_model(model_name)
+            
+            if is_fine_tuned:
+                logger.info("Using fine-tuned model optimizations")
+                # Use special method for fine-tuned models
+                outputs = self.transcribe_audio_fine_tuned(
+                    audio_file, chunk_length_s, batch_size
+                )
+            else:
+                # Standard transcription for regular models
+                logger.info("Using standard model transcription")
+                
+                # Prepare generation kwargs
+                generate_kwargs = {}
+                
+                # Language handling
+                if language != "Automatic Detection" and not model_name.endswith(".en"):
+                    language_map = {
+                        "Greek": "greek",
+                        "English": "english",
+                        "Spanish": "spanish", 
+                        "French": "french",
+                        "German": "german",
+                        "Italian": "italian"
+                    }
+                    lang_code = language_map.get(language, language.lower())
+                    generate_kwargs["language"] = lang_code
+                    logger.info(f"Set language: {lang_code}")
+                
+                # Task handling
+                if not model_name.endswith(".en"):
+                    generate_kwargs["task"] = task
+                
                 outputs = self.pipe(
                     audio_file,
                     chunk_length_s=chunk_length_s,
@@ -208,43 +289,35 @@ class OptimizedWhisperApp:
                     generate_kwargs=generate_kwargs,
                     return_timestamps=return_timestamps,
                 )
-                logger.info("Transcription completed")
-            except Exception as transcribe_error:
-                logger.error(f"Transcription failed: {transcribe_error}")
-                return f"Transcription failed: {str(transcribe_error)}", "", ""
             
             transcription_time = time.time() - start_time
-            logger.info(f"Total processing time: {transcription_time:.2f} seconds")
+            logger.info(f"Transcription completed in {transcription_time:.2f} seconds")
             
-            # Extract and validate results
+            # Extract results
             transcription = outputs.get("text", "") if outputs else ""
             chunks = outputs.get("chunks", []) if outputs else []
             
-            logger.info(f"Extracted transcription length: {len(transcription)} chars")
-            logger.info(f"Number of chunks: {len(chunks)}")
-            
-            # Handle timestamps safely
+            # Handle timestamps
             timestamp_text = ""
             if return_timestamps:
                 try:
                     if chunks:
                         timestamp_text = self._format_timestamps(chunks)
                     else:
-                        timestamp_text = "=== TIMESTAMPS ===\nNo chunks returned by the model.\n"
+                        timestamp_text = "=== TIMESTAMPS ===\nNo chunks returned.\n"
                 except Exception as ts_error:
-                    logger.warning(f"Error formatting timestamps: {ts_error}")
-                    timestamp_text = f"=== TIMESTAMPS ===\nError formatting timestamps: {str(ts_error)}\n"
+                    logger.warning(f"Timestamp formatting error: {ts_error}")
+                    timestamp_text = f"=== TIMESTAMPS ===\nError: {str(ts_error)}\n"
             else:
-                timestamp_text = "=== TIMESTAMPS ===\nTimestamp output disabled.\n"
+                timestamp_text = "=== TIMESTAMPS ===\nDisabled.\n"
             
             # Create detailed output
             detailed_output = self._format_detailed_output(
                 transcription, model_name, language, task, 
                 transcription_time, chunk_length_s, batch_size,
-                use_flash_attention, len(chunks)
+                use_flash_attention, len(chunks), is_fine_tuned
             )
             
-            logger.info("=== Transcription completed successfully ===")
             return transcription.strip(), timestamp_text, detailed_output
             
         except Exception as e:
@@ -255,14 +328,11 @@ class OptimizedWhisperApp:
             return error_msg, "", error_msg
     
     def _format_timestamps(self, chunks):
-        """Format timestamp information with comprehensive error handling"""
+        """Format timestamp information"""
         timestamp_text = "=== TIMESTAMPS ===\n"
         
         if not chunks:
-            timestamp_text += "No timestamp information available.\n"
-            return timestamp_text
-        
-        logger.info(f"Formatting {len(chunks)} chunks")
+            return timestamp_text + "No chunks available.\n"
         
         for i, chunk in enumerate(chunks):
             try:
@@ -274,32 +344,31 @@ class OptimizedWhisperApp:
                 elif isinstance(timestamp, (list, tuple)) and len(timestamp) >= 2:
                     start, end = timestamp[0], timestamp[1]
                     if start is None or end is None:
-                        timestamp_text += f"[Invalid timestamp]: {text}\n"
+                        timestamp_text += f"[Invalid]: {text}\n"
                     else:
                         try:
-                            start_f = float(start) if start is not None else 0.0
-                            end_f = float(end) if end is not None else 0.0
+                            start_f = float(start)
+                            end_f = float(end)
                             timestamp_text += f"[{start_f:.1f}s - {end_f:.1f}s]: {text}\n"
-                        except (ValueError, TypeError) as ve:
+                        except (ValueError, TypeError):
                             timestamp_text += f"[Format error]: {text}\n"
                 else:
                     timestamp_text += f"[Unexpected format]: {text}\n"
-                    
-            except Exception as chunk_error:
-                logger.warning(f"Error processing chunk {i}: {chunk_error}")
-                timestamp_text += f"[Chunk {i} error]: Processing failed\n"
+            except Exception as e:
+                timestamp_text += f"[Chunk {i} error]: {str(e)}\n"
         
         return timestamp_text
     
     def _format_detailed_output(self, transcription, model_name, language, task, 
                                transcription_time, chunk_length_s, batch_size, 
-                               use_flash_attention, num_chunks):
+                               use_flash_attention, num_chunks, is_fine_tuned=False):
         """Format detailed information"""
         output = "=== TRANSCRIPTION ===\n"
         output += f"{transcription}\n\n"
         
         output += "=== MODEL INFORMATION ===\n"
         output += f"Model: {model_name}\n"
+        output += f"Model Type: {'Fine-tuned' if is_fine_tuned else 'Standard'}\n"
         output += f"Language: {language}\n"
         output += f"Task: {task}\n"
         output += f"Processing time: {transcription_time:.2f} seconds\n"
@@ -308,26 +377,15 @@ class OptimizedWhisperApp:
         output += "\n=== PROCESSING SETTINGS ===\n"
         output += f"Chunk length: {chunk_length_s} seconds\n"
         output += f"Batch size: {batch_size}\n"
-        output += f"Flash Attention: {'Enabled' if use_flash_attention else 'Disabled'}\n"
+        output += f"Flash Attention: {'Enabled' if use_flash_attention and not is_fine_tuned else 'Disabled'}\n"
         
-        if self.pipe and hasattr(self.pipe, 'model'):
-            try:
-                device = next(self.pipe.model.parameters()).device
-                dtype = next(self.pipe.model.parameters()).dtype
-                output += f"Device: {device}\n"
-                output += f"Data type: {dtype}\n"
-            except:
-                output += "Device info: Unable to retrieve\n"
-        
-        output += f"Flash Attention 2 available: {FLASH_ATTN_AVAILABLE and is_flash_attn_2_available()}\n"
-        
-        output += "\n=== SYSTEM STATUS ===\n"
-        if torch.cuda.is_available():
-            output += f"GPU Available: Yes\n"
-            output += f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB\n"
-            output += f"GPU Memory Used: {torch.cuda.memory_allocated() / 1e9:.1f}GB\n"
-        else:
-            output += "GPU Available: No\n"
+        if is_fine_tuned:
+            output += "\n=== FINE-TUNED MODEL OPTIMIZATIONS ===\n"
+            output += "• Conservative batch size (max 2)\n"
+            output += "• Float32 precision for stability\n"
+            output += "• Disabled flash attention\n"
+            output += "• Deterministic generation\n"
+            output += "• No beam search\n"
         
         return output
     
@@ -339,9 +397,10 @@ class OptimizedWhisperApp:
         try:
             device = next(self.pipe.model.parameters()).device
             dtype = next(self.pipe.model.parameters()).dtype
-            return f"✅ {self.current_model} loaded on {device} ({dtype})"
+            model_type = "Fine-tuned" if self.is_fine_tuned_model(self.current_model) else "Standard"
+            return f"✅ {self.current_model} ({model_type}) - {device} ({dtype})"
         except:
-            return f"✅ {self.current_model} loaded (device info unavailable)"
+            return f"✅ {self.current_model} loaded"
 
 # Initialize the app
 logger.info("Initializing Optimized Whisper App...")
@@ -349,9 +408,8 @@ whisper_app = OptimizedWhisperApp()
 
 def transcribe_wrapper(audio, model_name, language, task, chunk_length_s, 
                       batch_size, use_flash_attention, return_timestamps):
-    """Wrapper for Gradio interface with additional safety"""
+    """Wrapper for Gradio interface"""
     try:
-        logger.info(f"Transcribe wrapper called with model: {model_name}")
         return whisper_app.transcribe_audio(
             audio, model_name, language, task,
             chunk_length_s, batch_size, use_flash_attention, return_timestamps
@@ -365,6 +423,23 @@ def get_model_status():
     """Get current model status"""
     return whisper_app.get_model_info()
 
+def update_settings_for_model(model_name):
+    """Update recommended settings based on model type"""
+    is_fine_tuned = whisper_app.is_fine_tuned_model(model_name)
+    
+    if is_fine_tuned:
+        return {
+            "batch_size": gr.update(value=1, maximum=2),
+            "use_flash_attention": gr.update(value=False),
+            "chunk_length_s": gr.update(value=30)
+        }
+    else:
+        return {
+            "batch_size": gr.update(value=4, maximum=16),
+            "use_flash_attention": gr.update(value=False),
+            "chunk_length_s": gr.update(value=30)
+        }
+
 # Create the interface
 def create_interface():
     with gr.Blocks(title="Optimized Whisper Transcription", theme=gr.themes.Soft()) as interface:
@@ -373,13 +448,13 @@ def create_interface():
             """
             # 🚀 Optimized Whisper Transcription
             
-            **High-Performance Speech-to-Text with Enhanced Error Handling**
+            **Enhanced for Fine-tuned Models**
             
             Features:
-            - Comprehensive error handling and timeout protection
-            - Better logging for debugging
-            - Graceful handling of model loading issues
-            - Robust timestamp processing
+            - Special handling for fine-tuned models (like Greek dialect)
+            - Automatic optimization based on model type
+            - Conservative settings for stability
+            - Enhanced error handling
             """
         )
         
@@ -402,9 +477,9 @@ def create_interface():
                 # Model selection
                 model_dropdown = gr.Dropdown(
                     choices=whisper_app.available_models,
-                    value="openai/whisper-small",  # Start with smaller model
+                    value="openai/whisper-small",
                     label="Model",
-                    info="Start with 'small' model for testing"
+                    info="Auto-optimizes settings for fine-tuned models"
                 )
                 
                 # Basic settings
@@ -428,23 +503,22 @@ def create_interface():
                         maximum=60,
                         value=30,
                         step=5,
-                        label="Chunk Length (seconds)",
-                        info="30s is optimal for most cases"
+                        label="Chunk Length (seconds)"
                     )
                     
                     batch_size = gr.Slider(
                         minimum=1,
-                        maximum=8,  # Reduced default
-                        value=4,    # Smaller batch size
+                        maximum=16,
+                        value=4,
                         step=1,
                         label="Batch Size",
-                        info="Start with smaller batch size"
+                        info="Auto-adjusted for fine-tuned models"
                     )
                     
                     use_flash_attention = gr.Checkbox(
                         label="Flash Attention 2",
-                        value=False,  # Disabled by default for stability
-                        info="Enable only if you have compatible GPU and flash-attn installed"
+                        value=False,
+                        info="Auto-disabled for fine-tuned models"
                     )
                     
                     return_timestamps = gr.Checkbox(
@@ -489,28 +563,33 @@ def create_interface():
             show_progress=True
         )
         
-        # Update model status when model changes
+        # Auto-adjust settings when model changes
         model_dropdown.change(
-            fn=lambda: "Model will be loaded on next transcription",
-            outputs=[model_status]
+            fn=lambda model: (
+                f"Model will be loaded on next transcription ({'Fine-tuned' if whisper_app.is_fine_tuned_model(model) else 'Standard'} model)",
+                1 if whisper_app.is_fine_tuned_model(model) else 4,
+                False
+            ),
+            inputs=[model_dropdown],
+            outputs=[model_status, batch_size, use_flash_attention]
         )
         
         # Footer
         gr.Markdown(
             """
-            ### 🔧 Troubleshooting
-            
-            **If processing hangs:**
-            1. Start with `whisper-small` model
-            2. Disable Flash Attention
-            3. Use smaller batch size (1-4)
-            4. Check the console logs for errors
-            
-            **For best results:**
-            - Test with small model first
-            - Gradually increase model size if needed
-            - Monitor GPU memory usage
-            - Check logs for any error messages
+            ### 🎯 Fine-tuned Model Optimizations
+            
+            **Automatic optimizations for fine-tuned models:**
+            - Batch size limited to 1-2 for stability
+            - Flash Attention automatically disabled
+            - Float32 precision for better compatibility
+            - Conservative generation settings
+            - Enhanced error handling
+            
+            **For Greek dialect model specifically:**
+            - Use batch size 1
+            - Keep chunk length at 30 seconds
+            - Language detection usually works well
             """
         )
     
@@ -519,4 +598,4 @@ def create_interface():
 # Launch the app
 if __name__ == "__main__":
     interface = create_interface()
-    interface.launch(share=True, debug=True)
+    interface.launch(share=True)