Update app.py

app.py

@@ -1,12 +1,11 @@
 import os
-import pickle
 import numpy as np
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from transformers import (
-    AutoTokenizer, 
-    AutoModelForSeq2SeqLM, 
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
     AutoModelForTokenClassification,
     AutoModelForCausalLM,
     pipeline
@@ -17,12 +16,9 @@ from bs4 import BeautifulSoup
 import nltk
 import torch
 import pandas as pd
-import subprocess
-from typing import Dict, Optional
-import codecs
 from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file  # Import Safetensors loader
 
-    
 # Initialize FastAPI app
 app = FastAPI()
 
@@ -56,31 +52,31 @@ def load_models():
     """Initialize all required models"""
     try:
         print("Loading models...")
-        
+
         # Set device
         device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Device set to use {device}")
-        
+
         # Embedding models
         models['embedding'] = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
         models['cross_encoder'] = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512)
-        
+
         # Translation models
         models['ar_to_en_tokenizer'] = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
         models['ar_to_en_model'] = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
         models['en_to_ar_tokenizer'] = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
         models['en_to_ar_model'] = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
-        
+
         # NER model
         models['bio_tokenizer'] = AutoTokenizer.from_pretrained("blaze999/Medical-NER")
         models['bio_model'] = AutoModelForTokenClassification.from_pretrained("blaze999/Medical-NER")
         models['ner_pipeline'] = pipeline("ner", model=models['bio_model'], tokenizer=models['bio_tokenizer'])
-        
+
         # LLM model
         model_name = "M4-ai/Orca-2.0-Tau-1.8B"
         models['llm_tokenizer'] = AutoTokenizer.from_pretrained(model_name)
         models['llm_model'] = AutoModelForCausalLM.from_pretrained(model_name)
-        
+
         print("Models loaded successfully")
         return True
     except Exception as e:
@@ -89,40 +85,26 @@ def load_models():
 
 
 def load_embeddings() -> Optional[Dict[str, np.ndarray]]:
-    """Load embeddings from local file or HuggingFace Hub"""
+    """Load embeddings from Safetensors file"""
     try:
-        import pickle
-        import numpy as np
-        import os
-        from typing import Dict, Optional
-
-        embeddings_path = 'embeddings.pkl'
+        embeddings_path = 'embeddings.safetensors'
         if not os.path.exists(embeddings_path):
-            from huggingface_hub import hf_hub_download
             embeddings_path = hf_hub_download(
                 repo_id=os.environ.get('HF_SPACE_ID', ''),
-                filename="embeddings.pkl",
+                filename="embeddings.safetensors",
                 repo_type="space"
             )
 
-        class ASCIIUnpickler(pickle.Unpickler):
-            def find_class(self, module, name):
-                if module == "__main__":
-                    module = "numpy"
-                return super().find_class(module, name)
-
-        with open(embeddings_path, 'rb') as f:
-            unpickler = ASCIIUnpickler(f)
-            embeddings = unpickler.load()
-
+        embeddings = load_file(embeddings_path)
         if not isinstance(embeddings, dict):
-            return None
-            
-        return {k: np.array(v, dtype=np.float32) for k, v in embeddings.items()}
-            
+            raise ValueError("Invalid format for embeddings in Safetensors file.")
+
+        # Convert to dictionary with numpy arrays
+        return {k: tensor.numpy() for k, tensor in embeddings.items()}
     except Exception as e:
         print(f"Error loading embeddings: {e}")
         return None
+
         
 def load_documents_data():
     """Load document data with error handling"""