Files
semantic-kernel/python/semantic_kernel/kernel_config.py
T
Abby Harrison f6059cdf94 Python/Local Hugging Face Inference for Completions and Embeddings (#658)
### Motivation and Context
This PR introduces native python support for Hugging Face models that
can: complete text, generate new text, summarize, and that generate
embeddings. Currently only supports downloading the models locally from
the HF model hub. Future plans include supporting the HF inference API
as well.

### Description
- Added 2 services: `hf_text_completion` and `hf_text_embedding`
- `hf_text_completion` supports the following tasks: _text-generation_,
_text2text-generation_, and _summarization_
- `hf_text_embedding` supports any model supported by the
sentence-transformers pip package
- Added dependencies: pytorch, transformers, sentence-transformers to
`requirements.txt` and `poetry.lock`
- fixed typo: `get_embedding_service_service_id` ->
`get_embedding_service_id`
- Added a number of integration tests for supported HF models
2023-04-26 17:07:29 -07:00

261 lines
9.2 KiB
Python

# Copyright (c) Microsoft. All rights reserved.
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Type, TypeVar, Union
from semantic_kernel.connectors.ai.chat_completion_client_base import (
ChatCompletionClientBase,
)
from semantic_kernel.connectors.ai.embeddings.embedding_generator_base import (
EmbeddingGeneratorBase,
)
from semantic_kernel.connectors.ai.text_completion_client_base import (
TextCompletionClientBase,
)
from semantic_kernel.reliability.pass_through_without_retry import (
PassThroughWithoutRetry,
)
from semantic_kernel.reliability.retry_mechanism import RetryMechanism
if TYPE_CHECKING:
from semantic_kernel.kernel_base import KernelBase
T = TypeVar("T")
class KernelConfig:
def __init__(self) -> None:
self._text_services: Dict[
str, Callable[["KernelBase"], TextCompletionClientBase]
] = {}
self._chat_services: Dict[
str, Callable[["KernelBase"], ChatCompletionClientBase]
] = {}
self._embedding_services: Dict[
str, Callable[["KernelBase"], EmbeddingGeneratorBase]
] = {}
self._default_text_service: Optional[str] = None
self._default_chat_service: Optional[str] = None
self._default_embedding_service: Optional[str] = None
self._retry_mechanism: RetryMechanism = PassThroughWithoutRetry()
def get_ai_service(
self, type: Type[T], service_id: Optional[str] = None
) -> Callable[["KernelBase"], T]:
matching_type = {}
if type == TextCompletionClientBase:
service_id = service_id or self._default_text_service
matching_type = self._text_services
elif type == ChatCompletionClientBase:
service_id = service_id or self._default_chat_service
matching_type = self._chat_services
elif type == EmbeddingGeneratorBase:
service_id = service_id or self._default_embedding_service
matching_type = self._embedding_services
else:
raise ValueError(f"Unknown AI service type: {type.__name__}")
if service_id not in matching_type:
raise ValueError(
f"{type.__name__} service with service_id '{service_id}' not found"
)
return matching_type[service_id]
def all_text_services(self) -> List[str]:
return list(self._text_services.keys())
def all_chat_services(self) -> List[str]:
return list(self._chat_services.keys())
def all_embedding_services(self) -> List[str]:
return list(self._embedding_services.keys())
def add_text_service(
self,
service_id: str,
service: Union[
TextCompletionClientBase, Callable[["KernelBase"], TextCompletionClientBase]
],
overwrite: bool = True,
) -> "KernelConfig":
if not service_id:
raise ValueError("service_id must be a non-empty string")
if not overwrite and service_id in self._text_services:
raise ValueError(
f"Text service with service_id '{service_id}' already exists"
)
self._text_services[service_id] = (
service if isinstance(service, Callable) else lambda _: service
)
if self._default_text_service is None:
self._default_text_service = service_id
return self
def add_chat_service(
self,
service_id: str,
service: Union[
ChatCompletionClientBase, Callable[["KernelBase"], ChatCompletionClientBase]
],
overwrite: bool = True,
) -> "KernelConfig":
if not service_id:
raise ValueError("service_id must be a non-empty string")
if not overwrite and service_id in self._chat_services:
raise ValueError(
f"Chat service with service_id '{service_id}' already exists"
)
self._chat_services[service_id] = (
service if isinstance(service, Callable) else lambda _: service
)
if self._default_chat_service is None:
self._default_chat_service = service_id
if isinstance(service, TextCompletionClientBase):
self.add_text_service(service_id, service)
if self._default_text_service is None:
self._default_text_service = service_id
return self
def add_embedding_service(
self,
service_id: str,
service: Union[
EmbeddingGeneratorBase, Callable[["KernelBase"], EmbeddingGeneratorBase]
],
overwrite: bool = False,
) -> "KernelConfig":
if not service_id:
raise ValueError("service_id must be a non-empty string")
if not overwrite and service_id in self._embedding_services:
raise ValueError(
f"Embedding service with service_id '{service_id}' already exists"
)
self._embedding_services[service_id] = (
service if isinstance(service, Callable) else lambda _: service
)
if self._default_embedding_service is None:
self._default_embedding_service = service_id
return self
# TODO: look harder at retry stuff
def set_default_text_service(self, service_id: str) -> "KernelConfig":
if service_id not in self._text_services:
raise ValueError(
f"AI service with service_id '{service_id}' does not exist"
)
self._default_text_service = service_id
return self
def set_default_chat_service(self, service_id: str) -> "KernelConfig":
if service_id not in self._chat_services:
raise ValueError(
f"AI service with service_id '{service_id}' does not exist"
)
self._default_chat_service = service_id
return self
def set_default_embedding_service(self, service_id: str) -> "KernelConfig":
if service_id not in self._embedding_services:
raise ValueError(
f"AI service with service_id '{service_id}' does not exist"
)
self._default_embedding_service = service_id
return self
def get_text_service_service_id(self, service_id: Optional[str] = None) -> str:
if service_id is None or service_id not in self._text_services:
if self._default_text_service is None:
raise ValueError("No default text service is set")
return self._default_text_service
return service_id
def get_chat_service_service_id(self, service_id: Optional[str] = None) -> str:
if service_id is None or service_id not in self._chat_services:
if self._default_chat_service is None:
raise ValueError("No default chat service is set")
return self._default_chat_service
return service_id
def get_embedding_service_id(self, service_id: Optional[str] = None) -> str:
if service_id is None or service_id not in self._embedding_services:
if self._default_embedding_service is None:
raise ValueError("No default embedding service is set")
return self._default_embedding_service
return service_id
def remove_text_service(self, service_id: str) -> "KernelConfig":
if service_id not in self._text_services:
raise ValueError(
f"AI service with service_id '{service_id}' does not exist"
)
del self._text_services[service_id]
if self._default_text_service == service_id:
self._default_text_service = next(iter(self._text_services), None)
return self
def remove_chat_service(self, service_id: str) -> "KernelConfig":
if service_id not in self._chat_services:
raise ValueError(
f"AI service with service_id '{service_id}' does not exist"
)
del self._chat_services[service_id]
if self._default_chat_service == service_id:
self._default_chat_service = next(iter(self._chat_services), None)
return self
def remove_embedding_service(self, service_id: str) -> "KernelConfig":
if service_id not in self._embedding_services:
raise ValueError(
f"AI service with service_id '{service_id}' does not exist"
)
del self._embedding_services[service_id]
if self._default_embedding_service == service_id:
self._default_embedding_service = next(iter(self._embedding_services), None)
return self
def clear_all_text_services(self) -> "KernelConfig":
self._text_services = {}
self._default_text_service = None
return self
def clear_all_chat_services(self) -> "KernelConfig":
self._chat_services = {}
self._default_chat_service = None
return self
def clear_all_embedding_services(self) -> "KernelConfig":
self._embedding_services = {}
self._default_embedding_service = None
return self
def clear_all_services(self) -> "KernelConfig":
self._text_services = {}
self._chat_services = {}
self._embedding_services = {}
self._default_text_service = None
self._default_chat_service = None
self._default_embedding_service = None
return self