Chaitanya Sagar Gurujula commited on
Commit
f275595
·
1 Parent(s): 4810fdb

added initial version

Browse files
Files changed (1) hide show
  1. src/app.py +4 -6
src/app.py CHANGED
@@ -11,6 +11,7 @@ from fastapi.middleware.cors import CORSMiddleware
11
  from fastapi.responses import HTMLResponse
12
  from fastapi.staticfiles import StaticFiles
13
  from pathlib import Path
 
14
 
15
  # Get the absolute path to the templates directory
16
  TEMPLATES_DIR = os.path.join(os.path.dirname(__file__), "templates")
@@ -35,18 +36,15 @@ app.add_middleware(
35
  # Set device
36
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
37
 
38
- # Create cache directory in the current working directory
39
- cache_dir = Path("model_cache")
40
- cache_dir.mkdir(exist_ok=True)
41
-
42
- # Set environment variable for Hugging Face cache
43
  os.environ['TRANSFORMERS_CACHE'] = str(cache_dir)
44
  os.environ['HF_HOME'] = str(cache_dir)
45
 
46
  # Load model from Hugging Face Hub
47
  def load_model():
48
  try:
49
- # Download the model file from HF Hub with custom cache directory
50
  model_path = hf_hub_download(
51
  repo_id=MODEL_ID,
52
  filename="best_model.pth",
 
11
  from fastapi.responses import HTMLResponse
12
  from fastapi.staticfiles import StaticFiles
13
  from pathlib import Path
14
+ import tempfile
15
 
16
  # Get the absolute path to the templates directory
17
  TEMPLATES_DIR = os.path.join(os.path.dirname(__file__), "templates")
 
36
  # Set device
37
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
38
 
39
+ # Use system's temporary directory
40
+ cache_dir = Path(tempfile.gettempdir()) / "model_cache"
 
 
 
41
  os.environ['TRANSFORMERS_CACHE'] = str(cache_dir)
42
  os.environ['HF_HOME'] = str(cache_dir)
43
 
44
  # Load model from Hugging Face Hub
45
  def load_model():
46
  try:
47
+ # Download the model file from HF Hub with temp directory
48
  model_path = hf_hub_download(
49
  repo_id=MODEL_ID,
50
  filename="best_model.pth",