Upload 3 files

Dockerfile

@@ -10,6 +10,12 @@ COPY . /app
 # Install any needed packages specified in requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
 
+# Create a writable cache directory
+RUN mkdir -p /app/cache
+
+# Set environment variable for the Transformers cache directory
+ENV TRANSFORMERS_CACHE=/app/cache
+
 # Make port 80 available to the world outside this container
 EXPOSE 80
 

app.py

@@ -1,9 +1,13 @@
 from fastapi import FastAPI
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import os
 
 app = FastAPI()
 
+# Set the cache directory environment variable
+os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
+
 # Load the model and tokenizer
 model_name = "unsloth/llama-3-8b-bnb-4bit"
 tokenizer = AutoTokenizer.from_pretrained(model_name)