Spaces:
Runtime error
Runtime error
File size: 3,145 Bytes
a22e84b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import json
from typing import Any, Dict, Optional
from loguru import logger
try:
import boto3
except ModuleNotFoundError:
logger.warning("Couldn't load AWS or SageMaker imports. Run 'poetry install --with aws' to support AWS.")
from llm_engineering.domain.inference import Inference
from llm_engineering.settings import settings
class LLMInferenceSagemakerEndpoint(Inference):
"""
Class for performing inference using a SageMaker endpoint for LLM schemas.
"""
def __init__(
self,
endpoint_name: str,
default_payload: Optional[Dict[str, Any]] = None,
inference_component_name: Optional[str] = None,
) -> None:
super().__init__()
self.client = boto3.client(
"sagemaker-runtime",
region_name=settings.AWS_REGION,
aws_access_key_id=settings.AWS_ACCESS_KEY,
aws_secret_access_key=settings.AWS_SECRET_KEY,
)
self.endpoint_name = endpoint_name
self.payload = default_payload if default_payload else self._default_payload()
self.inference_component_name = inference_component_name
def _default_payload(self) -> Dict[str, Any]:
"""
Generates the default payload for the inference request.
Returns:
dict: The default payload.
"""
return {
"inputs": "How is the weather?",
"parameters": {
"max_new_tokens": settings.MAX_NEW_TOKENS_INFERENCE,
"top_p": settings.TOP_P_INFERENCE,
"temperature": settings.TEMPERATURE_INFERENCE,
"return_full_text": False,
},
}
def set_payload(self, inputs: str, parameters: Optional[Dict[str, Any]] = None) -> None:
"""
Sets the payload for the inference request.
Args:
inputs (str): The input text for the inference.
parameters (dict, optional): Additional parameters for the inference. Defaults to None.
"""
self.payload["inputs"] = inputs
if parameters:
self.payload["parameters"].update(parameters)
def inference(self) -> Dict[str, Any]:
"""
Performs the inference request using the SageMaker endpoint.
Returns:
dict: The response from the inference request.
Raises:
Exception: If an error occurs during the inference request.
"""
try:
logger.info("Inference request sent.")
invoke_args = {
"EndpointName": self.endpoint_name,
"ContentType": "application/json",
"Body": json.dumps(self.payload),
}
if self.inference_component_name not in ["None", None]:
invoke_args["InferenceComponentName"] = self.inference_component_name
response = self.client.invoke_endpoint(**invoke_args)
response_body = response["Body"].read().decode("utf8")
return json.loads(response_body)
except Exception:
logger.exception("SageMaker inference failed.")
raise
|