KeyArgo
/
private-gpt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Use OpenAI for embeddings when openai mode is selected (#1096)

This commit is contained in:
Iván Martínez 2023-10-23 10:50:42 +02:00 committed by GitHub
parent 769a047b54
commit 78546524d0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
private_gpt/components/embedding/embedding_component.py
View File

@ -13,14 +13,19 @@ class EmbeddingComponent:
@inject @inject
def __init__(self) -> None: def __init__(self) -> None:
match settings.llm.mode: match settings.llm.mode:
case "mock": case "local":
# Not a random number, is the dimensionality used by
# the default embedding model
self.embedding_model = MockEmbedding(384)
case _:
from llama_index.embeddings import HuggingFaceEmbedding from llama_index.embeddings import HuggingFaceEmbedding
self.embedding_model = HuggingFaceEmbedding( self.embedding_model = HuggingFaceEmbedding(
model_name=settings.local.embedding_hf_model_name, model_name=settings.local.embedding_hf_model_name,
cache_folder=str(models_cache_path), 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)