Update main.py

main.py

@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import uvicorn
@@ -7,15 +8,16 @@ from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 app = FastAPI()
 
 # Define the cache directory path within your project
-cache_dir = './cache'
+cache_dir = str(Path(__file__).parent.resolve() / 'cache')
 
 # Create the cache directory if it doesn't exist
-if not os.path.exists(cache_dir):
-    os.makedirs(cache_dir, exist_ok=True)
+os.makedirs(cache_dir, exist_ok=True)
 
 # Set the TRANSFORMERS_CACHE environment variable to the cache directory
 os.environ['TRANSFORMERS_CACHE'] = cache_dir
 
+print(f"Transformers cache directory: {os.environ['TRANSFORMERS_CACHE']}")
+
 # Ensure your Hugging Face token is set as an environment variable
 huggingface_token = os.environ.get("TOKEN")
 if not huggingface_token:
@@ -61,4 +63,4 @@ async def generate_text(item: Item):
         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)