mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-29 11:12:23 +00:00
87 lines
3.1 KiB
Python
87 lines
3.1 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
from typing import List, Optional, Union, AsyncIterator
|
|
|
|
from llama_stack.apis.common.content_types import InterleavedContent, InterleavedContentItem
|
|
from llama_stack.apis.inference import Inference, Message, ToolChoice, ResponseFormat, LogProbConfig, ToolConfig, \
|
|
ChatCompletionResponse, ChatCompletionResponseStreamChunk, EmbeddingsResponse, TextTruncation, EmbeddingTaskType
|
|
from llama_stack.models.llama.datatypes import SamplingParams, ToolDefinition, ToolPromptFormat
|
|
from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
|
|
|
|
from . import WatsonXConfig
|
|
|
|
from ibm_watson_machine_learning.foundation_models import Model
|
|
from ibm_watson_machine_learning.metanames import GenTextParamsMetaNames as GenParams
|
|
|
|
from .models import MODEL_ENTRIES
|
|
|
|
|
|
|
|
class WatsonXInferenceAdapter(Inference, ModelRegistryHelper):
|
|
def __init__(self, config: WatsonXConfig) -> None:
|
|
ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
|
|
|
|
print(f"Initializing WatsonXInferenceAdapter({config.url})...")
|
|
|
|
self._config = config
|
|
self._credential = {
|
|
"url": self._config.url,
|
|
"apikey": self._config.api_key
|
|
}
|
|
|
|
self._project_id = self._config.project_id
|
|
self.params = {
|
|
GenParams.MAX_NEW_TOKENS: 4096,
|
|
GenParams.STOP_SEQUENCES: ["<|endoftext|>"]
|
|
}
|
|
|
|
async def completion(
|
|
self,
|
|
model_id: str,
|
|
content: InterleavedContent,
|
|
sampling_params: Optional[SamplingParams] = None,
|
|
response_format: Optional[ResponseFormat] = None,
|
|
stream: Optional[bool] = False,
|
|
logprobs: Optional[LogProbConfig] = None,
|
|
):
|
|
pass
|
|
|
|
async def embeddings(
|
|
self,
|
|
model_id: str,
|
|
contents: List[str] | List[InterleavedContentItem],
|
|
text_truncation: Optional[TextTruncation] = TextTruncation.none,
|
|
output_dimension: Optional[int] = None,
|
|
task_type: Optional[EmbeddingTaskType] = None,
|
|
) -> EmbeddingsResponse:
|
|
pass
|
|
|
|
async def chat_completion(
|
|
self,
|
|
model_id: str,
|
|
messages: List[Message],
|
|
sampling_params: Optional[SamplingParams] = None,
|
|
response_format: Optional[ResponseFormat] = None,
|
|
tools: Optional[List[ToolDefinition]] = None,
|
|
tool_choice: Optional[ToolChoice] = ToolChoice.auto,
|
|
tool_prompt_format: Optional[ToolPromptFormat] = None,
|
|
stream: Optional[bool] = False,
|
|
logprobs: Optional[LogProbConfig] = None,
|
|
tool_config: Optional[ToolConfig] = None,
|
|
):
|
|
# Language model
|
|
model = Model(
|
|
model_id=model_id,
|
|
credentials=self._credential,
|
|
project_id=self._project_id,
|
|
)
|
|
prompt = "\n".join(messages) + "\nAI: "
|
|
|
|
response = model.generate_text(prompt=prompt, params=self.params)
|
|
|
|
return response
|
|
|