Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Remove global torch import in llm_inference_utils.py #11115

Merged
merged 3 commits into from Feb 14, 2024
Merged

Remove global torch import in llm_inference_utils.py #11115

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6640f60
Remove global torch import in llm_inference_utils.py
B-Step62 Feb 14, 2024
f843267
Dummy commit to trigger Langchain CI test
B-Step62 Feb 14, 2024
025b2f9
Revert dummy commit
B-Step62 Feb 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
8 changes: 6 additions & 2 deletions mlflow/transformers/llm_inference_utils.py
View file
Open in desktop
@@ -1,9 +1,10 @@
from __future__ import annotations

import time
import uuid
from typing import Any, Dict, List, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import pandas as pd
import torch
from transformers import AutoTokenizer, StoppingCriteria

from mlflow.exceptions import MlflowException
Expand All @@ -15,6 +16,9 @@
COMPLETIONS_MODEL_OUTPUT_SCHEMA,
)

if TYPE_CHECKING:
import torch

_LLM_INFERENCE_TASK_KEY = "inference_task"
# The LLM inference task is saved as "task" in the metadata for forward compatibility with
# future Databricks Provisioned Throughput support of more model architectures for inference.
Expand Down