Updated transcription models

Dockerfile

@@ -10,6 +10,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
+# Create cache directories with proper permissions
+RUN mkdir -p /tmp/transformers_cache /tmp/hf_home /tmp/xdg_cache /tmp/nltk_data \
+    && chmod 777 /tmp/transformers_cache /tmp/hf_home /tmp/xdg_cache /tmp/nltk_data
+
 # Copy requirements first for better caching
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
@@ -17,17 +21,17 @@ RUN pip install --no-cache-dir -r requirements.txt
 # Install additional dependencies directly
 RUN pip install --no-cache-dir sentence-transformers>=2.2.2
 
-# Download model files - only use the sentence transformer
-RUN python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('all-MiniLM-L6-v2')"
-
-# Skip Whisper model preloading - it will be downloaded on first use
-# We don't need to preload it during build
+# Setup environment variables
+ENV TRANSFORMERS_CACHE="/tmp/transformers_cache"
+ENV HF_HOME="/tmp/hf_home"
+ENV XDG_CACHE_HOME="/tmp/xdg_cache"
+ENV NLTK_DATA="/tmp/nltk_data"
 
-# Run NLTK downloads - add all needed NLTK data
-RUN python -c "import nltk; nltk.download('punkt'); nltk.download('wordnet'); nltk.download('omw-1.4')"
+# Download sentence transformer model to cache
+RUN python -c "from sentence_transformers import SentenceTransformer; model = SentenceTransformer('all-MiniLM-L6-v2'); print(f'Model loaded successfully: {model}')"
 
-# Cache the torch hub directory to avoid unnecessary downloads
-RUN mkdir -p /root/.cache/torch/hub
+# Run NLTK downloads to the tmp directory
+RUN python -c "import nltk; nltk.data.path.append('/tmp/nltk_data'); nltk.download('punkt', download_dir='/tmp/nltk_data'); nltk.download('wordnet', download_dir='/tmp/nltk_data'); nltk.download('omw-1.4', download_dir='/tmp/nltk_data'); print('NLTK data downloaded successfully')"
 
 # Copy the rest of the application
 COPY . .

fastapi_app.py

@@ -8,7 +8,7 @@ import base64
 import numpy as np
 import torch
 import time
-from transformers import pipeline
+from transformers import pipeline, AutoProcessor, AutoModelForSpeechSeq2Seq
 import threading
 import json
 import logging
@@ -16,11 +16,32 @@ from pydantic import BaseModel
 from typing import Optional, Dict, Any, List
 from topic_segmenter import TopicSegmenter
 from middleware import HTTPSProxyMiddleware
+import tempfile
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Set environment variables for cache directories to writable locations
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
+os.environ["HF_HOME"] = "/tmp/hf_home"
+os.environ["XDG_CACHE_HOME"] = "/tmp/xdg_cache"
+os.environ["NLTK_DATA"] = "/tmp/nltk_data"
+
+# Create directories
+for directory in ["/tmp/transformers_cache", "/tmp/hf_home", "/tmp/xdg_cache", "/tmp/nltk_data"]:
+    os.makedirs(directory, exist_ok=True)
+
+# Try to download NLTK data to the tmp directory
+try:
+    import nltk
+    nltk.data.path.append("/tmp/nltk_data")
+    nltk.download('punkt', download_dir="/tmp/nltk_data", quiet=True)
+    nltk.download('wordnet', download_dir="/tmp/nltk_data", quiet=True)
+    nltk.download('omw-1.4', download_dir="/tmp/nltk_data", quiet=True)
+except Exception as e:
+    logger.warning(f"NLTK download issue: {str(e)}")
+
 app = FastAPI(title="Flowify", description="Audio transcription and topic segmentation API")
 
 # Add custom middleware for HTTPS handling
@@ -74,50 +95,82 @@ def get_model(model_name):
         if model_name not in models:
             logging.info(f"Loading model: {model_name}")
             
-            hf_model_name = model_name
+            # Handle Xenova models differently from OpenAI models
             if model_name.startswith('Xenova/'):
-                base_name = model_name.split('/')[-1] 
-                if '.en' in base_name:
-                    size = base_name.split('.')[0].replace('whisper-', '')
-                    hf_model_name = f"openai/whisper-{size}"
-                else:
-                    size = base_name.replace('whisper-', '')
-                    hf_model_name = f"openai/whisper-{size}"
-                
-                logging.info(f"Converting Xenova model {model_name} to {hf_model_name}")
-            
-            try:
-                # Simple direct approach without from_tf
-                from transformers import pipeline
-                
-                models[model_name] = pipeline(
-                    "automatic-speech-recognition", 
-                    hf_model_name,
-                    chunk_length_s=30,
-                    stride_length_s=5
-                )
-                logging.info(f"Model {model_name} loaded successfully using direct pipeline")
-            except Exception as e:
-                logging.error(f"Error loading model with direct pipeline: {str(e)}")
                 try:
-                    # Try with explicit loading but no from_tf
-                    from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, pipeline
+                    # For Xenova models, we'll use a different approach
+                    # We need to map to the corresponding HF model but handle it differently
+                    base_name = model_name.split('/')[-1]
+                    
+                    # Extract the size from the model name
+                    if '.en' in base_name:
+                        size = base_name.split('.')[0].replace('whisper-', '')
+                        hf_model_name = f"openai/whisper-{size}"
+                    else:
+                        size = base_name.replace('whisper-', '')
+                        hf_model_name = f"openai/whisper-{size}"
                     
-                    processor = AutoProcessor.from_pretrained(hf_model_name)
-                    model = AutoModelForSpeechSeq2Seq.from_pretrained(hf_model_name)
+                    logging.info(f"Using Xenova format for {model_name}, loading from {hf_model_name}")
                     
+                    # Load processor and model specifically for Whisper
+                    processor = AutoProcessor.from_pretrained(
+                        hf_model_name,
+                        cache_dir="/tmp/transformers_cache",
+                        local_files_only=False
+                    )
+                    
+                    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                        hf_model_name,
+                        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                        cache_dir="/tmp/transformers_cache",
+                        local_files_only=False
+                    )
+                    
+                    # Create pipeline with the loaded model and processor
                     models[model_name] = pipeline(
-                        "automatic-speech-recognition", 
+                        "automatic-speech-recognition",
                         model=model,
-                        tokenizer=processor,
-                        feature_extractor=processor,
+                        tokenizer=processor.tokenizer,
+                        feature_extractor=processor.feature_extractor,
                         chunk_length_s=30,
                         stride_length_s=5
                     )
-                    logging.info(f"Model {model_name} loaded successfully with explicit loading")
-                except Exception as e2:
-                    logging.error(f"Failed to load model {model_name}: {str(e2)}")
-                    raise RuntimeError(f"Failed to load model {model_name}: {str(e2)}")
+                    logging.info(f"Model {model_name} loaded successfully")
+                    
+                except Exception as e:
+                    logging.error(f"Error loading Xenova model: {str(e)}")
+                    raise RuntimeError(f"Failed to load Xenova model {model_name}: {str(e)}")
+            else:
+                # This is not a Xenova model, so we'll use the original approach
+                # but ensure we're using the cache directory
+                try:
+                    processor = AutoProcessor.from_pretrained(
+                        model_name,
+                        cache_dir="/tmp/transformers_cache",
+                        local_files_only=False
+                    )
+                    
+                    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                        model_name,
+                        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                        cache_dir="/tmp/transformers_cache",
+                        local_files_only=False
+                    )
+                    
+                    models[model_name] = pipeline(
+                        "automatic-speech-recognition",
+                        model=model,
+                        tokenizer=processor.tokenizer,
+                        feature_extractor=processor.feature_extractor,
+                        chunk_length_s=30,
+                        stride_length_s=5
+                    )
+                    logging.info(f"Model {model_name} loaded successfully")
+                    
+                except Exception as e:
+                    logging.error(f"Failed to load model {model_name}: {str(e)}")
+                    raise RuntimeError(f"Failed to load model {model_name}: {str(e)}")
+                
         return models[model_name]
 
 def process_audio(audio_data, sample_rate=16000, model_name="openai/whisper-base"):
@@ -192,7 +245,8 @@ async def analyze_topics(data: TranscriptData):
             min_segment_size=2,  
             topic_similarity_threshold=0.25,  
             max_topics=6,  
-            hierarchical_threshold=0.6  
+            hierarchical_threshold=0.6,
+            model_name="all-MiniLM-L6-v2"
         )
         
         segments, topic_mappings, topic_history, topic_hierarchies = segmenter.segment_transcript(transcript)

templates/index.html

@@ -59,9 +59,9 @@
                                 <div class="select-wrapper">
                                     <select id="modelSelector">
                                         <option value="Xenova/whisper-tiny.en">Whisper Tiny (Fast)</option>
-                                        <option value="Xenova/whisper-base.en">Whisper Base</option>
-                                        <option value="Xenova/whisper-small.en" selected>Whisper Small</option>
-                                        <option value="Xenova/whisper-medium.en">Whisper Medium</option>
+                                        <option value="Xenova/whisper-base.en" selected>Whisper Base</option>
+                                        <option value="Xenova/whisper-small.en">Whisper Small</option>
+                                        <option value="Xenova/whisper-medium.en">Whisper Medium (Slow)</option>
                                     </select>
                                     <i class="fas fa-chevron-down select-icon"></i>
                                 </div>

topic_segmenter.py

@@ -1,22 +1,68 @@
 import numpy as np
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
-from nltk.tokenize import sent_tokenize, word_tokenize
-from nltk.stem import WordNetLemmatizer
-from keybert import KeyBERT
 import re
 from collections import defaultdict
 from sklearn.cluster import AgglomerativeClustering
 import logging
+import os
+
+# Set NLTK data path to tmp directory
+nltk_data_dir = "/tmp/nltk_data"
+os.makedirs(nltk_data_dir, exist_ok=True)
+os.environ["NLTK_DATA"] = nltk_data_dir
 
 # Ensure NLTK resources are available
 try:
     import nltk
-    nltk.download('punkt', quiet=True)
-    nltk.download('wordnet', quiet=True)
-    nltk.download('omw-1.4', quiet=True)
+    nltk.data.path.append(nltk_data_dir)
+    
+    # Download resources if needed
+    for resource in ['punkt', 'wordnet', 'omw-1.4']:
+        try:
+            nltk.download(resource, download_dir=nltk_data_dir, quiet=True)
+        except Exception as e:
+            logging.warning(f"Error downloading NLTK resource {resource}: {str(e)}")
+    
+    # Now import specific NLTK components
+    from nltk.tokenize import sent_tokenize, word_tokenize
+    from nltk.stem import WordNetLemmatizer
 except Exception as e:
-    logging.warning(f"NLTK resource download issue: {str(e)}")
+    logging.warning(f"NLTK import/setup issue: {str(e)}")
+    # Fallback implementations if NLTK fails
+    def sent_tokenize(text):
+        # Simple fallback sentence tokenizer
+        return text.replace('\n', ' ').split('. ')
+    
+    def word_tokenize(text):
+        # Simple fallback word tokenizer
+        return text.split()
+    
+    class WordNetLemmatizer:
+        def lemmatize(self, word):
+            return word
+
+# Try to import KeyBERT, but provide fallback if it fails
+try:
+    from keybert import KeyBERT
+except ImportError as e:
+    logging.warning(f"KeyBERT import error: {str(e)}")
+    # Define a simple fallback KeyBERT class
+    class KeyBERT:
+        def __init__(self, model=None):
+            self.model = model
+        
+        def extract_keywords(self, text, **kwargs):
+            # Simple fallback keyword extraction using frequency
+            if not text:
+                return [("topic", 1.0)]
+            words = text.lower().split()
+            word_count = defaultdict(int)
+            for word in words:
+                if len(word) > 3:  # Skip very short words
+                    word_count[word] += 1
+            sorted_words = sorted(word_count.items(), key=lambda x: x[1], reverse=True)
+            return sorted_words[:kwargs.get('top_n', 5)]
 
 class TopicSegmenter:
     def __init__(self, window_size=3, similarity_threshold=0.2, context_size=2, min_segment_size=2, topic_similarity_threshold=0.3, 
@@ -34,16 +80,20 @@ class TopicSegmenter:
             import os
             os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Prevent parallelism warning
             
+            # Set cache directories
+            os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
+            os.environ["HF_HOME"] = "/tmp/hf_home"
+            
             # Try using sentence-transformers directly first
             try:
                 from sentence_transformers import SentenceTransformer
-                embedding_model = SentenceTransformer(model_name)
+                embedding_model = SentenceTransformer(model_name, cache_folder="/tmp/transformers_cache")
                 self.keyword_extractor = KeyBERT(model=embedding_model)
             except (ImportError, ValueError, RuntimeError) as e:
                 # Fallback to default model with lower-level settings
+                logging.warning(f"Using fallback model for KeyBERT: {str(e)}")
                 self.keyword_extractor = KeyBERT(model='distilbert-base-nli-mean-tokens')
         except Exception as e:
-            import logging
             logging.error(f"Error initializing KeyBERT model: {str(e)}")
             # Last resort fallback to simple keyword extraction without a model
             self.keyword_extractor = None