47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
"""
|
|
The ``mlflow.groq`` module provides an API for logging and loading Groq models.
|
|
"""
|
|
|
|
from mlflow.groq._groq_autolog import patched_call
|
|
from mlflow.telemetry.events import AutologgingEvent
|
|
from mlflow.telemetry.track import _record_event
|
|
from mlflow.utils.autologging_utils import autologging_integration, safe_patch
|
|
|
|
FLAVOR_NAME = "groq"
|
|
|
|
|
|
@autologging_integration(FLAVOR_NAME)
|
|
def autolog(
|
|
log_traces: bool = True,
|
|
disable: bool = False,
|
|
silent: bool = False,
|
|
):
|
|
"""
|
|
Enables (or disables) and configures autologging from Groq to MLflow.
|
|
Only synchronous calls are supported. Asynchronous APIs and streaming are not recorded.
|
|
|
|
Args:
|
|
log_traces: If ``True``, traces are logged for Groq models. If ``False``, no traces are
|
|
collected during inference. Default to ``True``.
|
|
disable: If ``True``, disables the Groq autologging. Default to ``False``.
|
|
silent: If ``True``, suppress all event logs and warnings from MLflow during Groq
|
|
autologging. If ``False``, show all events and warnings.
|
|
"""
|
|
|
|
from groq.resources.audio.transcriptions import Transcriptions
|
|
from groq.resources.audio.translations import Translations
|
|
from groq.resources.chat.completions import Completions as ChatCompletions
|
|
from groq.resources.embeddings import Embeddings
|
|
|
|
for task in (ChatCompletions, Translations, Transcriptions, Embeddings):
|
|
safe_patch(
|
|
FLAVOR_NAME,
|
|
task,
|
|
"create",
|
|
patched_call,
|
|
)
|
|
|
|
_record_event(
|
|
AutologgingEvent, {"flavor": FLAVOR_NAME, "log_traces": log_traces, "disable": disable}
|
|
)
|