vhr1007 commited on
Commit
88a7359
·
1 Parent(s): 298c89d
Files changed (2) hide show
  1. Dockerfile +3 -0
  2. utils/embedding_utils.py +8 -1
Dockerfile CHANGED
@@ -4,6 +4,9 @@ FROM python:3.10.9
4
  # Set the working directory to /app
5
  WORKDIR /app
6
 
 
 
 
7
  # Create the cache directory
8
  #RUN mkdir -p /app/.cache/huggingface
9
 
 
4
  # Set the working directory to /app
5
  WORKDIR /app
6
 
7
+ # Set a writable directory for Hugging Face cache
8
+ # ENV TRANSFORMERS_CACHE=/app/huggingface_cache
9
+
10
  # Create the cache directory
11
  #RUN mkdir -p /app/.cache/huggingface
12
 
utils/embedding_utils.py CHANGED
@@ -17,6 +17,13 @@ import logging
17
  # Load environment variables from .env file
18
  #load_dotenv()
19
 
 
 
 
 
 
 
 
20
  # Load Hugging Face token from environment variable
21
  # huggingface_token = os.getenv('HUGGINGFACE_HUB_TOKEN')
22
 
@@ -27,7 +34,7 @@ import logging
27
  # raise ValueError("Hugging Face token is not set. Please set the HUGGINGFACE_HUB_TOKEN environment variable.")
28
 
29
  # Load model and tokenizer with authentication
30
- #model_name = 'nvidia/NV-Embed-v1'
31
 
32
  model_name = 'nomic-ai/nomic-embed-text-v1.5'
33
  model = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5', trust_remote_code=True)
 
17
  # Load environment variables from .env file
18
  #load_dotenv()
19
 
20
+ os.environ["TRANSFORMERS_CACHE"] = "/app/huggingface_cache"
21
+
22
+ # Ensure the cache directory exists
23
+ cache_dir = os.environ["TRANSFORMERS_CACHE"]
24
+ if not os.path.exists(cache_dir):
25
+ os.makedirs(cache_dir)
26
+
27
  # Load Hugging Face token from environment variable
28
  # huggingface_token = os.getenv('HUGGINGFACE_HUB_TOKEN')
29
 
 
34
  # raise ValueError("Hugging Face token is not set. Please set the HUGGINGFACE_HUB_TOKEN environment variable.")
35
 
36
  # Load model and tokenizer with authentication
37
+ model_name = 'nvidia/NV-Embed-v1'
38
 
39
  model_name = 'nomic-ai/nomic-embed-text-v1.5'
40
  model = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5', trust_remote_code=True)