Update app.py

app.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import pipeline
 import logging
+import os
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -9,12 +10,17 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI(title="SQL Coder API")
 
+# Ensure cache directory exists
+cache_dir = os.getenv('TRANSFORMERS_CACHE', '/home/user/.cache/huggingface')
+os.makedirs(cache_dir, exist_ok=True)
+
 # Initialize pipeline
 try:
     pipe = pipeline("text-generation", 
                    model="defog/llama-3-sqlcoder-8b",
                    device_map="auto",
-                   torch_dtype="auto")
+                   torch_dtype="auto",
+                   cache_dir=cache_dir)
     logger.info("Pipeline initialized successfully")
 except Exception as e:
     logger.error(f"Error initializing pipeline: {str(e)}")