Update main.py

main.py

@@ -1,58 +1,50 @@
 import os
-
-# Hugging Face Spaces: redirect cache to /tmp
-os.environ["HF_HUB_CACHE"] = "/tmp/hf_cache"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
-
-import os
-
-# Redirect Hugging Face cache to local project directory (not /)
-os.environ["HF_HOME"] = "./hf_cache"
-os.environ["TRANSFORMERS_CACHE"] = "./hf_cache"
-os.environ["HF_DATASETS_CACHE"] = "./hf_cache"
-os.environ["HF_METRICS_CACHE"] = "./hf_cache"
-
+import logging
 import uvicorn
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from huggingface_hub import snapshot_download
 from app.api.routes import router
 from app.core.config import settings
-import logging
-import os
-from huggingface_hub import snapshot_download
 
+# ====== Step 1: Safe Hugging Face cache location ======
+os.environ["HF_HOME"] = "./hf_cache"
+os.environ["TRANSFORMERS_CACHE"] = "./hf_cache"
+os.environ["HF_HUB_CACHE"] = "./hf_cache"
+os.environ["HF_DATASETS_CACHE"] = "./hf_cache"
+os.environ["HF_METRICS_CACHE"] = "./hf_cache"
+
+# ====== Step 2: Logging setup ======
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-
-# Hugging Face model & dataset download
 logger.info("Downloading model and dataset from Hugging Face Hub...")
 
-HF_TOKEN = os.getenv("HF_TOKEN")  # Set in environment variables
+# ====== Step 3: Load your Hugging Face token ======
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Download FAISS index dataset
-FAISS_INDEX_PATH = snapshot_download(
-    repo_id="negi2725/dataRag",
-    repo_type="dataset",
+# ====== Step 4: Download model and FAISS index ======
+MODEL_PATH = snapshot_download(
+    repo_id="negi2725/legalBert",
     token=HF_TOKEN,
-    local_dir="./hf_cache/faiss_index",
+    local_dir="./models/legalbert_model",
     local_dir_use_symlinks=False
 )
 
-# Download LegalBERT model
-MODEL_PATH = snapshot_download(
-    repo_id="negi2725/legalBert",
+FAISS_INDEX_PATH = snapshot_download(
+    repo_id="negi2725/dataRag",
     token=HF_TOKEN,
-    local_dir="./hf_cache/legalbert_model",
+    local_dir="./faiss_indexes",
     local_dir_use_symlinks=False
 )
 
 logger.info(f"FAISS index files downloaded to: {FAISS_INDEX_PATH}")
 logger.info(f"Model files downloaded to: {MODEL_PATH}")
 
-# Make these paths accessible globally in your app if needed
-os.environ["FAISS_INDEX_PATH"] = FAISS_INDEX_PATH
+# ====== Step 5: Make paths globally available ======
 os.environ["MODEL_PATH"] = MODEL_PATH
+os.environ["FAISS_INDEX_PATH"] = FAISS_INDEX_PATH
 
+# ====== Step 6: Initialize FastAPI ======
 app = FastAPI(
     title="Legal RAG Analysis API",
     description="FastAPI backend for legal case analysis using RAG system with LegalBERT predictions and Gemini AI evaluation",
@@ -67,6 +59,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Register API routes
 app.include_router(router, prefix="/api/v1")
 
 @app.get("/")
@@ -77,6 +70,7 @@ async def root():
 async def health_check():
     return {"status": "healthy", "message": "API is running"}
 
+# ====== Step 7: Run locally if needed ======
 if __name__ == "__main__":
     uvicorn.run(
         "main:app",