Update app.py

app.py

@@ -19,8 +19,48 @@ from transformers import (
 import pandas as pd
 import time
 
+# Define persistent_load to handle persistent IDs
+def persistent_load(pers_id):
+    """
+    Handle persistent IDs during unpickling.
+    """
+    if pers_id == "sentence_transformer_model":
+        try:
+            # Load a pre-defined SentenceTransformer model
+            model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+            return model
+        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}")
+
+# Function to load models
+def load_models():
+    try:
+        # Load embeddings data with custom persistent_load function
+        with open("models/embeddings.pkl", "rb") as file:
+            global_models.embeddings_data = pickle.load(file, persistent_load=persistent_load)
+            print("Embeddings data loaded successfully.")
+    except pickle.UnpicklingError as e:
+        raise HTTPException(status_code=500, detail=f"Unpickling error: {e}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to load models: {e}")
+        
 app = FastAPI()
 
+@app.on_event("startup")
+async def startup_event():
+    """
+    Load models at application startup.
+    """
+    print("Loading models...")
+    load_models()
+    print("Models loaded.")
+
+@app.get("/")
+async def root():
+    return {"message": "Server is running"}
+    
 # Models and data structures to store loaded models
 class GlobalModels:
     embedding_model = None