Update main.py

main.py

@@ -1,24 +1,31 @@
+import os
+
+# Set cache directories to writable locations
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface_cache"
+os.environ["HF_HOME"] = "/tmp/hf_home"  
+os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/huggingface_cache"
+
+os.makedirs("/tmp/huggingface_cache", exist_ok=True)
+os.makedirs("/tmp/hf_home", exist_ok=True)
+
 import pdfplumber
 import re
 import nltk
 import torch
 import uvicorn
-import os
 import time
 from nltk.tokenize import sent_tokenize
 from transformers import pipeline
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 
-# ✅ Set Hugging Face Cache Directory to /tmp/ (Fix Permission Issues)
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface_cache"
+# ✅ Set Hugging Face Token for Authentication
+from huggingface_hub import login
 
 # ✅ Ensure NLTK Dependencies are Stored in a Writable Directory
 NLTK_DATA_DIR = "/tmp/nltk_data"  
 os.makedirs(NLTK_DATA_DIR, exist_ok=True)
 nltk.data.path.append(NLTK_DATA_DIR)
-
-# ✅ Download punkt tokenizer
 try:
     nltk.data.find("tokenizers/punkt")
 except LookupError:
@@ -40,7 +47,7 @@ app.add_middleware(
 device = 0 if torch.cuda.is_available() else -1
 print(f"Using Device: {'GPU' if device == 0 else 'CPU'}")
 
-# ✅ Load Summarization Model (Fixing Cache Issue)
+# ✅ Load Summarization Model
 summarizer = pipeline("summarization", model="google/pegasus-xsum", device=device)
 
 # --- **Generalized Cleaning** ---
@@ -95,6 +102,5 @@ async def summarize_pdf(file: UploadFile = File(...)):
     except Exception as e:
         return {"error": str(e)}
 
-# ✅ Run FastAPI Server
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)