mirror of
https://github.com/imartinez/privateGPT.git
synced 2025-04-29 20:13:33 +00:00
895588b82a
* fix: docker copying extra files * feat: allow configuring mode through env vars * feat: Attempt to build and tag a docker image * fix: run docker on release * fix: typing in prompt transformation * chore: remove tutorial comments
32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
from injector import inject, singleton
|
|
from llama_index import MockEmbedding
|
|
from llama_index.embeddings.base import BaseEmbedding
|
|
|
|
from private_gpt.paths import models_cache_path
|
|
from private_gpt.settings.settings import settings
|
|
|
|
|
|
@singleton
|
|
class EmbeddingComponent:
|
|
embedding_model: BaseEmbedding
|
|
|
|
@inject
|
|
def __init__(self) -> None:
|
|
match settings.llm.mode:
|
|
case "local" | "sagemaker":
|
|
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 "openai":
|
|
from llama_index import OpenAIEmbedding
|
|
|
|
openai_settings = settings.openai.api_key
|
|
self.embedding_model = OpenAIEmbedding(api_key=openai_settings)
|
|
case "mock":
|
|
# Not a random number, is the dimensionality used by
|
|
# the default embedding model
|
|
self.embedding_model = MockEmbedding(384)
|