updates

ai_med_extract/api/routes.py

@@ -107,15 +107,36 @@ def run_qa_pipeline(qa_pipeline, question, context):
         raise
 
 def get_ner_pipeline(ner_model_type, ner_model_name):
+    if not ner_model_type or not ner_model_name:
+        raise ValueError("Both ner_model_type and ner_model_name must be provided")
+
     if not hasattr(get_ner_pipeline, "cache"):
         get_ner_pipeline.cache = {}
+
     key = (ner_model_type, ner_model_name)
     if key not in get_ner_pipeline.cache:
-        from transformers import pipeline
+        try:
+            from transformers import pipeline
+            logging.info(f"Loading NER pipeline - Type: {ner_model_type}, Model: {ner_model_name}")
+
+            get_ner_pipeline.cache[key] = pipeline(
+                task=ner_model_type,
+                model=ner_model_name,
+                trust_remote_code=True,
+                device_map="auto"
+            )
+            logging.info(f"Successfully loaded NER pipeline for {ner_model_name}")
+        except Exception as e:
+            logging.error(f"Failed to load NER pipeline: {str(e)}", exc_info=True)
+            if "Connection" in str(e):
+                raise RuntimeError(f"Network error while loading model: {str(e)}")
+            elif "CUDA" in str(e):
+                raise RuntimeError(f"GPU error while loading model: {str(e)}")
+            elif "disk space" in str(e):
+                raise RuntimeError(f"Insufficient disk space: {str(e)}")
+            else:
+                raise RuntimeError(f"Error loading model: {str(e)}")
 
-        get_ner_pipeline.cache[key] = pipeline(
-            task=ner_model_type, model=ner_model_name, trust_remote_code=True
-        )
     return get_ner_pipeline.cache[key]
 
 
@@ -201,34 +222,45 @@ def register_routes(app, agents):
             try:
                 file.save(filepath)
             except Exception as e:
-                return jsonify({"error": f"Failed to save file: {str(e)}"}), 500
+                return jsonify({"error": f"Filed to save file: {str(e)}"}), 500
             ext = filename.rsplit(".", 1)[-1].lower()
-            extracted_text = TextExtractorAgent.extract_text(filepath, ext)
-            if not extracted_text or extracted_text == "No text found":
-                return (
-                    jsonify({"error": f"Failed to extract text from {filename}"}),
-                    415,
-                )
+            try:
+                extracted_text = TextExtractorAgent.extract_text(filepath, ext)
+                if not extracted_text or extracted_text == "No text found":
+                    os.remove(filepath)  # Clean up on failure
+                    return (
+                        jsonify({"error": f"Failed to extract text from {filename}"}),
+                        415,
+                    )
+            except Exception as e:
+                logging.error(f"Text extraction failed for {filename}: {str(e)}", exc_info=True)
+                os.remove(filepath)  # Clean up on failure
+                return jsonify({"error": f"Text extraction failed: {str(e)}"}), 500
+
             skip_medical_check = (
                 request.form.get("skip_medical_check", "false").lower() == "true"
             )
             if not skip_medical_check:
-                ner_results = ner_pipeline(extracted_text)
-                medical_entities = list(
-                    set(
-                        [
-                            r["word"]
-                            for r in ner_results
-                            if r["entity"].startswith("B-")
-                            or r["entity"].startswith("I-")
-                        ]
-                    )
-                )
-                if not medical_entities:
-                    return (
-                        jsonify({"error": f"'{filename}' is not medically relevant"}),
-                        406,
+                try:
+                    ner_results = ner_pipeline(extracted_text)
+                    medical_entities = list(
+                        set(
+                            [
+                                r["word"]
+                                for r in ner_results
+                                if r["entity"].startswith("B-")
+                                or r["entity"].startswith("I-")
+                            ]
+                        )
                     )
+                    if not medical_entities:
+                        return (
+                            jsonify({"error": f"'{filename}' is not medically relevant"}),
+                            406,
+                        )
+                except Exception as e:
+                    logging.error(f"NER processing failed for {filename}: {str(e)}", exc_info=True)
+                    return jsonify({"error": f"NER processing failed: {str(e)}"}), 500
             skip_patient_check = (
                 request.form.get("skip_patient_check", "false").lower() == "true"
             )
@@ -260,7 +292,8 @@ def register_routes(app, agents):
                 else:
                     summary = str(summary_result)
             except Exception as e:
-                summary = "Summary failed"
+                logging.error(f"Summary generation failed for {filename}: {str(e)}", exc_info=True)
+                summary = f"Summary generation failed: {str(e)}"
             extracted_data.append(
                 {
                     "file": filename,
@@ -378,37 +411,98 @@ def register_routes(app, agents):
 
     @app.route("/api/extract_medical_data_from_audio", methods=["POST"])
     def extract_medical_data_from_audio():
-        if "audio" not in request.files:
-            return jsonify({"error": "No audio file provided"}), 400
-        audio_file = request.files["audio"]
-        if audio_file.filename == "":
-            return jsonify({"error": "No selected audio file"}), 400
-        temp_path = os.path.join("/tmp", audio_file.filename)
-        audio_file.save(temp_path)
+        temp_path = None
         try:
-            result = whisper_model.transcribe(temp_path)
-            transcribed_text = result["text"]
+            # Validate request
+            if "audio" not in request.files:
+                return jsonify({"error": "No audio file provided"}), 400
+            
+            audio_file = request.files["audio"]
+            if audio_file.filename == "":
+                return jsonify({"error": "No selected audio file"}), 400
+                
+            # Validate file extension
+            if not allowed_file(audio_file.filename):
+                return jsonify({"error": f"Unsupported audio format. Allowed formats: wav, mp3, m4a, ogg"}), 400
+                
+            # Check file size
+            valid_size, error_message = check_file_size(audio_file)
+            if not valid_size:
+                return jsonify({"error": error_message}), 400
+                
+            # Generate a secure temporary path
+            import uuid
+            from werkzeug.utils import secure_filename
+            temp_filename = f"{uuid.uuid4()}_{secure_filename(audio_file.filename)}"
+            temp_path = os.path.join("/tmp", temp_filename)
+            
             try:
-                clean_text = PHIScrubberAgent.scrub_phi(transcribed_text)
-            except Exception:
-                clean_text = transcribed_text
-            summary = SummarizerAgent.generate_summary(clean_text)
-            medical_data = MedicalDataExtractorAgent.extract_medical_data(clean_text)
-            os.remove(temp_path)
-            return (
-                jsonify(
-                    {
-                        "transcribed_text": clean_text,
-                        "summary": summary,
-                        "medical_chart": medical_data,
-                    }
-                ),
-                200,
-            )
+                logging.info(f"Saving audio file to temporary path: {temp_path}")
+                audio_file.save(temp_path)
+                
+                # Initialize whisper model with retries
+                max_retries = 3
+                for attempt in range(max_retries):
+                    try:
+                        logging.info(f"Initializing Whisper model (attempt {attempt + 1}/{max_retries})")
+                        transcribed_text = whisper_model.transcribe(temp_path)["text"]
+                        if not transcribed_text:
+                            raise ValueError("No text output from transcription")
+                        logging.info("Audio transcription successful")
+                        break
+                    except Exception as e:
+                        if attempt == max_retries - 1:  # Last attempt
+                            raise
+                        logging.warning(f"Transcription attempt {attempt + 1} failed: {str(e)}")
+                        continue
+                
+                # Clean and process text
+                try:
+                    logging.info("Scrubbing PHI from transcribed text")
+                    clean_text = PHIScrubberAgent.scrub_phi(transcribed_text)
+                except Exception as e:
+                    logging.warning(f"PHI scrubbing failed, using raw text: {str(e)}")
+                    clean_text = transcribed_text
+                
+                try:
+                    logging.info("Generating summary")
+                    summary = SummarizerAgent.generate_summary(clean_text)
+                except Exception as e:
+                    logging.error(f"Summary generation failed: {str(e)}")
+                    summary = "Summary generation failed"
+                
+                try:
+                    logging.info("Extracting medical data")
+                    medical_data = MedicalDataExtractorAgent.extract_medical_data(clean_text)
+                except Exception as e:
+                    logging.error(f"Medical data extraction failed: {str(e)}")
+                    medical_data = {"error": f"Medical data extraction failed: {str(e)}"}
+                
+                # Clean up temporary file
+                if os.path.exists(temp_path):
+                    os.remove(temp_path)
+                
+                return jsonify({
+                    "transcribed_text": clean_text,
+                    "summary": summary,
+                    "medical_chart": medical_data,
+                }), 200
+                
+            except Exception as e:
+                logging.error(f"Audio processing failed: {str(e)}", exc_info=True)
+                if os.path.exists(temp_path):
+                    os.remove(temp_path)
+                return jsonify({
+                    "error": f"Audio processing failed: {str(e)}",
+                    "details": "Error occurred during audio transcription or text processing"
+                }), 500
+                
         except Exception as e:
-            if os.path.exists(temp_path):
-                os.remove(temp_path)
-            return jsonify({"error": f"Processing failed: {str(e)}"}), 500
+            logging.error(f"Request handling failed: {str(e)}", exc_info=True)
+            return jsonify({
+                "error": "Internal server error",
+                "details": str(e)
+            }), 500
 
     @app.route("/extract_medical_data_questions", methods=["POST"])
     def extract_medical_data_questions():

ai_med_extract/app.py

@@ -20,22 +20,67 @@ load_dotenv()
 app = Flask(__name__)
 CORS(app)
 
+# Configure upload directory
 UPLOAD_DIR = os.getenv('UPLOAD_DIR', os.path.join(os.getcwd(), 'uploads'))
-os.makedirs(UPLOAD_DIR, exist_ok=True)
+try:
+    os.makedirs(UPLOAD_DIR, exist_ok=True)
+    os.chmod(UPLOAD_DIR, 0o777)  # Ensure directory is writable
+except Exception as e:
+    logging.error(f"Failed to create/configure upload directory: {e}", exc_info=True)
+    UPLOAD_DIR = '/tmp/uploads'  # Fallback to /tmp if main directory creation fails
+    os.makedirs(UPLOAD_DIR, exist_ok=True)
+
 app.config['UPLOAD_FOLDER'] = UPLOAD_DIR
-app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16 MB max file size
+app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 * 1024  # 16 GB max file size to handle large medical files
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
-# Model loaders (example, adjust as needed)
-medalpaca_model_loader = None  # TODO: Implement LazyModelLoader if needed
-summarization_model_loader = None  # TODO: Implement LazyModelLoader if needed
-whisper_model = None
-def get_whisper_model():
-    global whisper_model
-    if whisper_model is None:
-        whisper_model = whisper.load_model("tiny")
-    return whisper_model
+# Model loaders
+class LazyModelLoader:
+    def __init__(self, model_name, model_type):
+        self.model_name = model_name
+        self.model_type = model_type
+        self._model = None
+
+    def load(self):
+        from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+        if self._model is None:
+            self._model = pipeline(
+                task=self.model_type,
+                model=self.model_name,
+                trust_remote_code=True,
+                device_map="auto"
+            )
+        return self._model
+
+medalpaca_model_loader = LazyModelLoader("medalpaca/medalpaca-13b", "text-generation")
+summarization_model_loader = LazyModelLoader("facebook/bart-large-cnn", "summarization")
+class WhisperModelLoader:
+    _instance = None
+    
+    def __init__(self):
+        self._model = None
+        
+    @staticmethod
+    def get_instance():
+        if WhisperModelLoader._instance is None:
+            WhisperModelLoader._instance = WhisperModelLoader()
+        return WhisperModelLoader._instance
+    
+    def load(self):
+        if self._model is None:
+            try:
+                logging.info("Loading Whisper model...")
+                self._model = whisper.load_model("base")
+                logging.info("Whisper model loaded successfully")
+            except Exception as e:
+                logging.error(f"Failed to load Whisper model: {str(e)}", exc_info=True)
+                raise RuntimeError(f"Failed to load Whisper model: {str(e)}")
+        return self._model
+        
+    def transcribe(self, audio_path):
+        model = self.load()
+        return model.transcribe(audio_path)
 
 # Initialize agents
 text_extractor_agent = TextExtractorAgent()
@@ -49,7 +94,7 @@ agents = {
     "phi_scrubber": phi_scrubber_agent,
     "summarizer": summarizer_agent,
     "medical_data_extractor": medical_data_extractor_agent,
-    "whisper_model": get_whisper_model
+    "whisper_model": WhisperModelLoader.get_instance()
 }
 
 from .api.routes import register_routes

ai_med_extract/utils/file_utils.py

@@ -5,7 +5,7 @@ import logging
 from werkzeug.utils import secure_filename
 from flask import current_app
 
-ALLOWED_EXTENSIONS = {"pdf", "jpg", "jpeg", "png", "svg", "docx", "doc", "xlsx", "xls"}
+ALLOWED_EXTENSIONS = {"pdf", "jpg", "jpeg", "png", "svg", "docx", "doc", "xlsx", "xls", "wav", "mp3", "m4a", "ogg"}
 MAX_SIZE_PDF_DOCS = 1 * 1024 * 1024 * 1024  # 1GB
 MAX_SIZE_IMAGES = 500 * 1024 * 1024  # 500MB
 
@@ -15,15 +15,26 @@ def allowed_file(filename):
 
 
 def check_file_size(file):
-    file.seek(0, os.SEEK_END)
-    size = file.tell()
-    file.seek(0)
-    extension = file.filename.rsplit('.', 1)[-1].lower()
-    if extension in {"pdf", "docx"} and size > MAX_SIZE_PDF_DOCS:
-        return False, f"File {file.filename} exceeds 1GB size limit"
-    elif extension in {"jpg", "jpeg", "png"} and size > MAX_SIZE_IMAGES:
-        return False, f"Image {file.filename} exceeds 500MB size limit"
-    return True, None
+    try:
+        # Store current position
+        current_pos = file.tell()
+        
+        # Check size
+        file.seek(0, os.SEEK_END)
+        size = file.tell()
+        
+        # Return to original position
+        file.seek(current_pos)
+        
+        extension = file.filename.rsplit('.', 1)[-1].lower()
+        if extension in {"pdf", "docx"} and size > MAX_SIZE_PDF_DOCS:
+            return False, f"File {file.filename} exceeds 1GB size limit"
+        elif extension in {"jpg", "jpeg", "png"} and size > MAX_SIZE_IMAGES:
+            return False, f"Image {file.filename} exceeds 500MB size limit"
+        return True, None
+    except Exception as e:
+        logging.error(f"Error checking file size: {e}", exc_info=True)
+        return False, f"Error checking file size: {str(e)}"
 
 
 def save_data_to_storage(filename, data):