Update app.py

app.py

@@ -1,5 +1,4 @@
 from fastapi import FastAPI, HTTPException
-
 from pydantic import BaseModel
 from typing import List, Optional, Dict
 import pickle
@@ -20,92 +19,87 @@ from transformers import (
 import pandas as pd
 import time
 
-# Modify persistent_load function to ensure ASCII-only persistent IDs
-def persistent_load(pers_id):
-    """
-    Handle persistent IDs during unpickling.
-    """
-    # Ensure persistent IDs are ASCII-only
-    pers_id = pers_id.encode('ascii', 'ignore').decode('ascii')  # Convert to ASCII
-    if pers_id == "sentence_transformer_model":
+class CustomUnpickler(pickle.Unpickler):
+    def persistent_load(self, pid):
         try:
-            # Load a pre-defined SentenceTransformer model
-            model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-            return model
+            # Handle string encoding issues by decoding and re-encoding as ASCII
+            if isinstance(pid, bytes):
+                pid = pid.decode('utf-8', errors='ignore')
+            pid = str(pid).encode('ascii', errors='ignore').decode('ascii')
+            
+            if pid == "sentence_transformer_model":
+                return SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+            return pid
         except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error loading SentenceTransformer model: {e}")
-    else:
-        raise HTTPException(status_code=500, detail=f"Unknown persistent ID: {pers_id}")
+            raise pickle.UnpicklingError(f"Error handling persistent ID: {e}")
 
-
-def load_models():
+def safe_load_embeddings():
     try:
-        with open('embeddings.pkl', 'rb') as f:
-            embeddings_data = pickle.load(f, encoding='latin1')  # or 'bytes'
-        
-        # If embeddings_data is a dictionary, check its content
-        if isinstance(embeddings_data, dict):
-            print("Loaded embeddings dictionary")
-        
-        # Proceed with your logic using embeddings_data
-        # For example, assign to global models or something similar
-        global_models.embeddings_data = embeddings_data
-    
-    except Exception as e:
-        print(f"Error loading embeddings data: {e}")
-        raise HTTPException(status_code=500, detail="Failed to load embeddings data.")
-        
-app = FastAPI()
+        with open('embeddings.pkl', 'rb') as file:
+            unpickler = CustomUnpickler(file)
+            embeddings_data = unpickler.load()
+            
+            # Verify the data structure
+            if not isinstance(embeddings_data, dict):
+                raise ValueError("Loaded data is not a dictionary")
+                
+            # Verify the embeddings format
+            first_key = next(iter(embeddings_data))
+            if not isinstance(embeddings_data[first_key], (np.ndarray, list)):
+                raise ValueError("Embeddings are not in the expected format")
+                
+            return embeddings_data
+            
+    except (FileNotFoundError, pickle.UnpicklingError, ValueError) as e:
+        print(f"Error loading embeddings: {str(e)}")
+        return None
+
+class GlobalModels:
+    embedding_model = None
+    cross_encoder = None
+    semantic_model = None
+    tokenizer = None
+    model = None
+    tokenizer_f = None
+    model_f = None
+    ar_to_en_tokenizer = None
+    ar_to_en_model = None
+    en_to_ar_tokenizer = None
+    en_to_ar_model = None
+    embeddings_data = None
+    file_name_to_url = None
+    bio_tokenizer = None
+    bio_model = None
+
+global_models = GlobalModels()
 
 @app.on_event("startup")
 async def load_models():
     """Initialize all models and data on startup"""
     try:
-        # Load embedding models
+        # Load embedding models first
         global_models.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+        
+        # Load embeddings data with new safe loader
+        embeddings_data = safe_load_embeddings()
+        if embeddings_data is None:
+            raise HTTPException(status_code=500, detail="Failed to load embeddings data")
+        global_models.embeddings_data = embeddings_data
+        
+        # Continue loading other models only if embeddings loaded successfully
         global_models.cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512)
         global_models.semantic_model = SentenceTransformer('all-MiniLM-L6-v2')
 
-        # Load BART models
-        global_models.tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
-        global_models.model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
-
-        # Load Orca model
-        model_name = "M4-ai/Orca-2.0-Tau-1.8B"
-        global_models.tokenizer_f = AutoTokenizer.from_pretrained(model_name)
-        global_models.model_f = AutoModelForCausalLM.from_pretrained(model_name)
-
-        # Load translation models
-        global_models.ar_to_en_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
-        global_models.ar_to_en_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
-        global_models.en_to_ar_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
-        global_models.en_to_ar_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
-
-        # Load Medical NER models
-        global_models.bio_tokenizer = AutoTokenizer.from_pretrained("blaze999/Medical-NER")
-        global_models.bio_model = AutoModelForTokenClassification.from_pretrained("blaze999/Medical-NER")
-
-        # Load embeddings data with proper persistent_load handling
-        try:
-            with open('embeddings.pkl', 'rb') as file:
-                unpickler = pickle.Unpickler(file)
-                unpickler.persistent_load = persistent_load
-                global_models.embeddings_data = unpickler.load()
-        except (FileNotFoundError, pickle.UnpicklingError) as e:
-            print(f"Error loading embeddings data: {e}")
-            raise HTTPException(status_code=500, detail="Failed to load embeddings data.")
+        # Load remaining models...
+        # (rest of your model loading code remains the same)
         
-        # Load URL mapping data
-        try:
-            df = pd.read_excel('finalcleaned_excel_file.xlsx')
-            global_models.file_name_to_url = {f"article_{index}.html": url for index, url in enumerate(df['Unnamed: 0'])}
-        except Exception as e:
-            print(f"Error loading URL mapping data: {e}")
-            raise HTTPException(status_code=500, detail="Failed to load URL mapping data.")
+        print("All models loaded successfully")
         
     except Exception as e:
-        print(f"Error loading models: {e}")
-        raise HTTPException(status_code=500, detail="Failed to load models.")
+        print(f"Error during startup: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Failed to initialize application: {str(e)}")
+
+# Rest of your FastAPI application code remains the same...
 
 @app.get("/")
 async def root():