Update app.py

app.py

@@ -8,8 +8,11 @@ app = FastAPI()
 @app.get("/")
 async def generate_text():
     try:
+        # Get the cache directory from the environment variable
+        cache_dir = os.getenv("TRANSFORMERS_CACHE", "/code/cache")
+
         # Specify the directory for model download within the Docker container
-        model_dir = "/code/model"
+        model_dir = os.path.join(cache_dir, "TheBloke/Mistral-7B-v0.1-GGUF")
         os.makedirs(model_dir, exist_ok=True)
         
         hf_hub_download(repo_id="TheBloke/Mistral-7B-v0.1-GGUF", filename="mistral-7b-v0.1.Q4_K_M.gguf", local_dir=model_dir)
@@ -23,4 +26,4 @@ async def generate_text():
         generated_texts = tokenizer.batch_decode(output, skip_special_tokens=True)
         return generated_texts
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e))