Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,896 Bytes
df6c67d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import os.path
import shutil
from typing import List, Optional, Union
from inference.core.env import MODEL_CACHE_DIR
from inference.core.utils.file_system import (
dump_bytes,
dump_json,
dump_text_lines,
read_json,
read_text_file,
)
def initialise_cache(model_id: Optional[str] = None) -> None:
cache_dir = get_cache_dir(model_id=model_id)
os.makedirs(cache_dir, exist_ok=True)
def are_all_files_cached(files: List[str], model_id: Optional[str] = None) -> bool:
return all(is_file_cached(file=file, model_id=model_id) for file in files)
def is_file_cached(file: str, model_id: Optional[str] = None) -> bool:
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
return os.path.isfile(cached_file_path)
def load_text_file_from_cache(
file: str,
model_id: Optional[str] = None,
split_lines: bool = False,
strip_white_chars: bool = False,
) -> Union[str, List[str]]:
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
return read_text_file(
path=cached_file_path,
split_lines=split_lines,
strip_white_chars=strip_white_chars,
)
def load_json_from_cache(
file: str, model_id: Optional[str] = None, **kwargs
) -> Optional[Union[dict, list]]:
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
return read_json(path=cached_file_path, **kwargs)
def save_bytes_in_cache(
content: bytes,
file: str,
model_id: Optional[str] = None,
allow_override: bool = True,
) -> None:
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
dump_bytes(path=cached_file_path, content=content, allow_override=allow_override)
def save_json_in_cache(
content: Union[dict, list],
file: str,
model_id: Optional[str] = None,
allow_override: bool = True,
**kwargs,
) -> None:
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
dump_json(
path=cached_file_path, content=content, allow_override=allow_override, **kwargs
)
def save_text_lines_in_cache(
content: List[str],
file: str,
model_id: Optional[str] = None,
allow_override: bool = True,
) -> None:
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
dump_text_lines(
path=cached_file_path, content=content, allow_override=allow_override
)
def get_cache_file_path(file: str, model_id: Optional[str] = None) -> str:
cache_dir = get_cache_dir(model_id=model_id)
return os.path.join(cache_dir, file)
def clear_cache(model_id: Optional[str] = None) -> None:
cache_dir = get_cache_dir(model_id=model_id)
if os.path.exists(cache_dir):
shutil.rmtree(cache_dir)
def get_cache_dir(model_id: Optional[str] = None) -> str:
if model_id is not None:
return os.path.join(MODEL_CACHE_DIR, model_id)
return MODEL_CACHE_DIR
|