SystemSh0cker/private-gpt
1
0
Fork 0
mirror of https://github.com/zylon-ai/private-gpt.git synced 2025-12-22 13:55:41 +01:00
Code Issues Projects Releases Packages Activity

Feature/sagemaker embedding (#1161)

Browse source 
* Sagemaker deployed embedding model support

---------

Co-authored-by: Pablo Orgaz <pabloogc@gmail.com>
This commit is contained in:
Iván Martínez 2023-11-05 16:16:49 +01:00 • committed by GitHub
parent f29df84301
commit ad512e3c42
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 114 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

11
private_gpt/components/embedding/embedding_component.py
View file

@ -13,13 +13,22 @@ class EmbeddingComponent:
@inject
def __init__(self) -> None:
match settings.llm.mode:
case "local" | "sagemaker":
case "local":
from llama_index.embeddings import HuggingFaceEmbedding
self.embedding_model = HuggingFaceEmbedding(
model_name=settings.local.embedding_hf_model_name,
cache_folder=str(models_cache_path),
)
case "sagemaker":
from private_gpt.components.embedding.custom.sagemaker import (
SagemakerEmbedding,
)
self.embedding_model = SagemakerEmbedding(
endpoint_name=settings.sagemaker.embedding_endpoint_name,
)
case "openai":
from llama_index import OpenAIEmbedding