Skip to content

Commit

Permalink
fix: filter spurious llama.cpp output
Browse files Browse the repository at this point in the history
  • Loading branch information
lsorber committed Dec 3, 2024
1 parent 405818e commit 463b54e
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion src/raglite/_litellm.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
"""Add support for llama-cpp-python models to LiteLLM."""

import asyncio
import contextlib
import logging
import warnings
from collections.abc import AsyncIterator, Callable, Iterator
from functools import cache
from io import StringIO
from typing import Any, ClassVar, cast

import httpx
Expand Down Expand Up @@ -96,7 +98,10 @@ def llm(model: str, **kwargs: Any) -> Llama:
filename, n_ctx_str = filename_n_ctx
n_ctx = int(n_ctx_str)
# Load the LLM.
with warnings.catch_warnings(): # Filter huggingface_hub warning about HF_TOKEN.
with (
contextlib.redirect_stdout(StringIO()), # Filter spurious llama.cpp output.
warnings.catch_warnings(), # Filter huggingface_hub warning about HF_TOKEN.
):
warnings.filterwarnings("ignore", category=UserWarning)
llm = Llama.from_pretrained(
repo_id=repo_id,
Expand Down

0 comments on commit 463b54e

Please sign in to comment.