diff --git a/pipelines/tools/batch_bill_summarizer.py b/pipelines/tools/batch_bill_summarizer.py index 33c5e6d..ee1a1c5 100644 --- a/pipelines/tools/batch_bill_summarizer.py +++ b/pipelines/tools/batch_bill_summarizer.py @@ -23,14 +23,10 @@ import httpx import typer from tiktoken import Encoding, get_encoding +from pipelines.config import get_config_dir from pipelines.tools.bill_token_compression import compress_bill_text -_PROMPTS_PATH = ( - Path(__file__).resolve().parents[2] - / "config" - / "prompts" - / "summarization_prompts.toml" -) +_PROMPTS_PATH = get_config_dir() / "prompts" / "summarization_prompts.toml" _PROMPTS = tomllib.loads(_PROMPTS_PATH.read_text())["summarization"] SUMMARIZATION_SYSTEM_PROMPT: str = _PROMPTS["system_prompt"] SUMMARIZATION_USER_TEMPLATE: str = _PROMPTS["user_template"] diff --git a/pipelines/tools/compresion_test.py b/pipelines/tools/compresion_test.py index c84c298..3814a20 100644 --- a/pipelines/tools/compresion_test.py +++ b/pipelines/tools/compresion_test.py @@ -24,14 +24,10 @@ from typing import Annotated import httpx import typer +from pipelines.config import get_config_dir from pipelines.tools.bill_token_compression import compress_bill_text -_PROMPTS_PATH = ( - Path(__file__).resolve().parents[2] - / "config" - / "prompts" - / "summarization_prompts.toml" -) +_PROMPTS_PATH = get_config_dir() / "prompts" / "summarization_prompts.toml" _PROMPTS = tomllib.loads(_PROMPTS_PATH.read_text())["summarization"] SUMMARIZATION_SYSTEM_PROMPT: str = _PROMPTS["system_prompt"] SUMMARIZATION_USER_TEMPLATE: str = _PROMPTS["user_template"] diff --git a/pipelines/tools/finetune.py b/pipelines/tools/finetune.py index 3fccf57..33df584 100644 --- a/pipelines/tools/finetune.py +++ b/pipelines/tools/finetune.py @@ -25,6 +25,8 @@ from datasets import Dataset from transformers import TrainingArguments from trl import SFTTrainer +from pipelines.config import default_config_path + logger = logging.getLogger(__name__) @@ -123,7 +125,7 @@ def main( config_path: Annotated[ Path, typer.Option("--config", help="TOML config file"), - ] = Path(__file__).parent / "config.toml", + ] = default_config_path(), save_gguf: Annotated[ bool, typer.Option("--save-gguf/--no-save-gguf", help="Also save GGUF") ] = False, diff --git a/pipelines/tools/prompt_bench.py b/pipelines/tools/prompt_bench.py index 60391c5..c126fc5 100644 --- a/pipelines/tools/prompt_bench.py +++ b/pipelines/tools/prompt_bench.py @@ -11,8 +11,8 @@ from typing import Annotated import typer -from pipelines.tools.containers.lib import check_gpu_free -from pipelines.tools.containers.vllm import start_vllm, stop_vllm +from pipelines.containers.lib import check_gpu_free +from pipelines.containers.vllm import start_vllm, stop_vllm from pipelines.tools.downloader import is_model_present from pipelines.tools.models import BenchmarkConfig from pipelines.tools.vllm_client import VLLMClient