Chittrarasu commited on
Commit
e95d03f
·
1 Parent(s): c84fd87
Files changed (1) hide show
  1. services/embedding.py +7 -9
services/embedding.py CHANGED
@@ -5,18 +5,16 @@ from typing import List, Optional
5
  import torch
6
  import os
7
 
8
- # Set multiple environment variables to redirect all caching
9
- os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
10
- os.environ["HF_HOME"] = "/app/cache"
11
- os.environ["XDG_CACHE_HOME"] = "/app/cache" # Additional variable for broader compatibility
12
 
13
- # Ensure cache directory exists (redundant with Dockerfile but added for safety)
14
- cache_dir = "/app/cache"
15
- if not os.path.exists(cache_dir):
16
- os.makedirs(cache_dir, exist_ok=True)
17
 
18
- # Load model with custom cache directory
19
  model = SentenceTransformer("clip-ViT-B-32")
 
 
20
 
21
  def get_text_embedding(text: str) -> Optional[List[float]]:
22
  try:
 
5
  import torch
6
  import os
7
 
8
+ from huggingface_hub import HfFileSystem
9
+ import transformers.utils.hub
 
 
10
 
11
+ # Set cache directory globally for huggingface_hub
12
+ transformers.utils.hub.HUGGINGFACE_HUB_CACHE = "/app/cache/hub"
 
 
13
 
14
+ # Load model
15
  model = SentenceTransformer("clip-ViT-B-32")
16
+ # Load model with custom cache directory
17
+ # model = SentenceTransformer("clip-ViT-B-32")
18
 
19
  def get_text_embedding(text: str) -> Optional[List[float]]:
20
  try: