Skip to content

Commit

Permalink
add hf cache cleanup fixture
Browse files Browse the repository at this point in the history
  • Loading branch information
Benjoyo committed Apr 29, 2024
1 parent 37411b2 commit 78ffd94
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions bpm-ai/tests/conftest.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
import os
import platform
import shutil

import pytest
from bpm_ai_core.llm.anthropic_chat.anthropic_chat import ChatAnthropic
from bpm_ai_core.llm.openai_chat.openai_chat import ChatOpenAI
from bpm_ai_core.util.rpc import remote_object
from bpm_ai_inference.util.hf import hf_home


@pytest.fixture
Expand All @@ -17,3 +20,11 @@ def local_llm():
@pytest.fixture
def llm():
return None #remote_object("ChatLlamaCpp", "0.0.0.0", 6666, model="QuantFactory/Phi-3-mini-4k-instruct-GGUF")


@pytest.fixture(autouse=True, scope="module")
def cleanup():
yield
if platform.system() != "Darwin":
cache_dir = hf_home()
shutil.rmtree(cache_dir, ignore_errors=True)

0 comments on commit 78ffd94

Please sign in to comment.