Skip to content
Snippets Groups Projects
Unverified Commit e861ccbe authored by Logan's avatar Logan Committed by GitHub
Browse files

Google genai embeddings (#18079)

parent ddcf5a39
No related branches found
No related tags found
No related merge requests found
Showing
with 748 additions and 2 deletions
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/gemini.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a> <a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/gemini.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
# Google Gemini Embeddings # Google Gemini Embeddings
**NOTE:** This example is deprecated. Please use the `GoogleGenAIEmbedding` class instead, detailed [here](https://github.com/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/google_genai.ipynb).
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙. If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
%pip install llama-index-embeddings-gemini %pip install llama-index-embeddings-gemini
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
!pip install llama-index 'google-generativeai>=0.3.0' matplotlib !pip install llama-index 'google-generativeai>=0.3.0' matplotlib
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
import os import os
GOOGLE_API_KEY = "" # add your GOOGLE API key here GOOGLE_API_KEY = "" # add your GOOGLE API key here
os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# imports # imports
from llama_index.embeddings.gemini import GeminiEmbedding from llama_index.embeddings.gemini import GeminiEmbedding
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# get API key and create embeddings # get API key and create embeddings
model_name = "models/embedding-001" model_name = "models/embedding-001"
embed_model = GeminiEmbedding( embed_model = GeminiEmbedding(
model_name=model_name, api_key=GOOGLE_API_KEY, title="this is a document" model_name=model_name, api_key=GOOGLE_API_KEY, title="this is a document"
) )
embeddings = embed_model.get_text_embedding("Google Gemini Embeddings.") embeddings = embed_model.get_text_embedding("Google Gemini Embeddings.")
``` ```
%% Output %% Output
/Users/haotianzhang/llama_index/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html /Users/haotianzhang/llama_index/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm from .autonotebook import tqdm as notebook_tqdm
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
print(f"Dimension of embeddings: {len(embeddings)}") print(f"Dimension of embeddings: {len(embeddings)}")
``` ```
%% Output %% Output
Dimension of embeddings: 768 Dimension of embeddings: 768
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
embeddings[:5] embeddings[:5]
``` ```
%% Output %% Output
[0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218] [0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218]
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
embeddings = embed_model.get_query_embedding("Google Gemini Embeddings.") embeddings = embed_model.get_query_embedding("Google Gemini Embeddings.")
embeddings[:5] embeddings[:5]
``` ```
%% Output %% Output
[0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218] [0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218]
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
embeddings = embed_model.get_text_embedding( embeddings = embed_model.get_text_embedding(
["Google Gemini Embeddings.", "Google is awesome."] ["Google Gemini Embeddings.", "Google is awesome."]
) )
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
print(f"Dimension of embeddings: {len(embeddings)}") print(f"Dimension of embeddings: {len(embeddings)}")
print(embeddings[0][:5]) print(embeddings[0][:5])
print(embeddings[1][:5]) print(embeddings[1][:5])
``` ```
%% Output %% Output
Dimension of embeddings: 2 Dimension of embeddings: 2
[0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218] [0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218]
[0.009427786, -0.009968997, -0.03341217, -0.025396815, 0.03210592] [0.009427786, -0.009968997, -0.03341217, -0.025396815, 0.03210592]
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
embedding = await embed_model.aget_text_embedding("Google Gemini Embeddings.") embedding = await embed_model.aget_text_embedding("Google Gemini Embeddings.")
print(embedding[:5]) print(embedding[:5])
``` ```
%% Output %% Output
[0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218] [0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218]
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
embeddings = await embed_model.aget_text_embedding_batch( embeddings = await embed_model.aget_text_embedding_batch(
[ [
"Google Gemini Embeddings.", "Google Gemini Embeddings.",
"Google is awesome.", "Google is awesome.",
"Llamaindex is awesome.", "Llamaindex is awesome.",
] ]
) )
print(embeddings[0][:5]) print(embeddings[0][:5])
print(embeddings[1][:5]) print(embeddings[1][:5])
print(embeddings[2][:5]) print(embeddings[2][:5])
``` ```
%% Output %% Output
[0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218] [0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218]
[0.009427786, -0.009968997, -0.03341217, -0.025396815, 0.03210592] [0.009427786, -0.009968997, -0.03341217, -0.025396815, 0.03210592]
[0.013159992, -0.021570021, -0.060150445, -0.042500723, 0.041159637] [0.013159992, -0.021570021, -0.060150445, -0.042500723, 0.041159637]
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
embedding = await embed_model.aget_query_embedding("Google Gemini Embeddings.") embedding = await embed_model.aget_query_embedding("Google Gemini Embeddings.")
print(embedding[:5]) print(embedding[:5])
``` ```
%% Output %% Output
[0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218] [0.028174246, -0.0290093, -0.013280814, 0.008629, 0.025442218]
......
%% Cell type:markdown id: tags:
<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/gemini.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
%% Cell type:markdown id: tags:
# Google GenAI Embeddings
Using Google's `google-genai` package, LlamaIndex provides a `GoogleGenAIEmbedding` class that allows you to embed text using Google's GenAI models from both the Gemini and Vertex AI APIs.
%% Cell type:markdown id: tags:
If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.
%% Cell type:code id: tags:
``` python
%pip install llama-index-embeddings-google-genai
```
%% Cell type:code id: tags:
``` python
import os
os.environ["GOOGLE_API_KEY"] = "..."
```
%% Cell type:markdown id: tags:
## Setup
`GoogleGenAIEmbedding` is a wrapper around the `google-genai` package, which means it supports both Gemini and Vertex AI APIs out of that box.
You can pass in the `api_key` directly, or pass in a `vertexai_config` to use the Vertex AI API.
Other options include `embed_batch_size`, `model_name`, and `embedding_config`.
The default model is `text-embedding-004`.
%% Cell type:code id: tags:
``` python
from llama_index.embeddings.google_genai import GoogleGenAIEmbedding
from google.genai.types import EmbedContentConfig
embed_model = GoogleGenAIEmbedding(
model_name="text-embedding-004",
embed_batch_size=100,
# can pass in the api key directly
# api_key="...",
# or pass in a vertexai_config
# vertexai_config={
# "project": "...",
# "location": "...",
# }
# can also pass in an embedding_config
# embedding_config=EmbedContentConfig(...)
)
```
%% Cell type:markdown id: tags:
## Usage
%% Cell type:markdown id: tags:
### Sync
%% Cell type:code id: tags:
``` python
embeddings = embed_model.get_text_embedding("Google Gemini Embeddings.")
print(embeddings[:5])
print(f"Dimension of embeddings: {len(embeddings)}")
```
%% Output
[0.031099992, 0.02192731, -0.06523498, 0.016788177, 0.0392835]
Dimension of embeddings: 768
%% Cell type:code id: tags:
``` python
embeddings = embed_model.get_query_embedding("Query Google Gemini Embeddings.")
print(embeddings[:5])
print(f"Dimension of embeddings: {len(embeddings)}")
```
%% Output
[0.022199392, 0.03671178, -0.06874573, 0.02195774, 0.05475164]
Dimension of embeddings: 768
%% Cell type:code id: tags:
``` python
embeddings = embed_model.get_text_embedding_batch(
[
"Google Gemini Embeddings.",
"Google is awesome.",
"Llamaindex is awesome.",
]
)
print(f"Got {len(embeddings)} embeddings")
print(f"Dimension of embeddings: {len(embeddings[0])}")
```
%% Output
Got 3 embeddings
Dimension of embeddings: 768
%% Cell type:markdown id: tags:
### Async
%% Cell type:code id: tags:
``` python
embeddings = await embed_model.aget_text_embedding("Google Gemini Embeddings.")
print(embeddings[:5])
print(f"Dimension of embeddings: {len(embeddings)}")
```
%% Output
[0.031099992, 0.02192731, -0.06523498, 0.016788177, 0.0392835]
Dimension of embeddings: 768
%% Cell type:code id: tags:
``` python
embeddings = await embed_model.aget_query_embedding(
"Query Google Gemini Embeddings."
)
print(embeddings[:5])
print(f"Dimension of embeddings: {len(embeddings)}")
```
%% Output
[0.022199392, 0.03671178, -0.06874573, 0.02195774, 0.05475164]
Dimension of embeddings: 768
%% Cell type:code id: tags:
``` python
embeddings = await embed_model.aget_text_embedding_batch(
[
"Google Gemini Embeddings.",
"Google is awesome.",
"Llamaindex is awesome.",
]
)
print(f"Got {len(embeddings)} embeddings")
print(f"Dimension of embeddings: {len(embeddings[0])}")
```
%% Output
Got 3 embeddings
Dimension of embeddings: 768
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/google_palm.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a> <a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/google_palm.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
# Google PaLM Embeddings # Google PaLM Embeddings
**NOTE:** This example is deprecated. Please use the `GoogleGenAIEmbedding` class instead, detailed [here](https://github.com/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/google_genai.ipynb).
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙. If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
%pip install llama-index-embeddings-google %pip install llama-index-embeddings-google
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
!pip install llama-index !pip install llama-index
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# imports # imports
from llama_index.embeddings.google import GooglePaLMEmbedding from llama_index.embeddings.google import GooglePaLMEmbedding
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# get API key and create embeddings # get API key and create embeddings
model_name = "models/embedding-gecko-001" model_name = "models/embedding-gecko-001"
api_key = "YOUR API KEY" api_key = "YOUR API KEY"
embed_model = GooglePaLMEmbedding(model_name=model_name, api_key=api_key) embed_model = GooglePaLMEmbedding(model_name=model_name, api_key=api_key)
embeddings = embed_model.get_text_embedding("Google PaLM Embeddings.") embeddings = embed_model.get_text_embedding("Google PaLM Embeddings.")
``` ```
%% Output %% Output
/opt/homebrew/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html /opt/homebrew/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm from .autonotebook import tqdm as notebook_tqdm
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
print(f"Dimension of embeddings: {len(embeddings)}") print(f"Dimension of embeddings: {len(embeddings)}")
``` ```
%% Output %% Output
Dimension of embeddings: 768 Dimension of embeddings: 768
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
embeddings[:5] embeddings[:5]
``` ```
%% Output %% Output
[0.028517298, -0.0028859433, -0.035110522, 0.021982985, -0.0039763353] [0.028517298, -0.0028859433, -0.035110522, 0.021982985, -0.0039763353]
......
# LlamaIndex Embeddings Integration: Gemini # LlamaIndex Embeddings Integration: Gemini
**NOTE:** This package is deprecated. Please use the `GoogleGenAIEmbedding` class instead, detailed [here](https://github.com/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/google_genai.ipynb).
llama_index/_static
.DS_Store
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
bin/
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
etc/
include/
lib/
lib64/
parts/
sdist/
share/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
.ruff_cache
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
notebooks/
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
pyvenv.cfg
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# Jetbrains
.idea
modules/
*.swp
# VsCode
.vscode
# pipenv
Pipfile
Pipfile.lock
# pyright
pyrightconfig.json
poetry_requirements(
name="poetry",
)
The MIT License
Copyright (c) Jerry Liu
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
help: ## Show all Makefile targets.
@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
format: ## Run code autoformatters (black).
pre-commit install
git ls-files | xargs pre-commit run black --files
lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy
pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
test: ## Run tests via pytest.
pytest tests
watch-docs: ## Build and watch documentation.
sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/
# Google GenAI Embeddings
This package provides a wrapper around the Google GenAI API, allowing you to use Gemini and Vertex AI embeddings in your projects.
## Installation
```bash
pip install llama-index-embeddings-google-genai
```
## Usage
```python
from llama_index.embeddings.google_genai import GoogleGenAIEmbedding
embed_model = GoogleGenAIEmbedding(model_name="text-embedding-004")
embeddings = embed_model.get_text_embedding("Hello, world!")
print(embeddings)
```
## Vertex AI
```python
embed_model = GoogleGenAIEmbedding(
model_name="text-embedding-004",
vertexai_config={
"project": "your-project-id",
"location": "your-location",
},
)
```
from llama_index.embeddings.google_genai.base import GoogleGenAIEmbedding
__all__ = ["GoogleGenAIEmbedding"]
"""Gemini embeddings file."""
import os
from typing import Any, Dict, List, Optional, TypedDict
from llama_index.core.base.embeddings.base import (
DEFAULT_EMBED_BATCH_SIZE,
BaseEmbedding,
)
from llama_index.core.bridge.pydantic import Field, PrivateAttr
from llama_index.core.callbacks.base import CallbackManager
import google.genai
import google.auth.credentials
import google.genai.types as types
class VertexAIConfig(TypedDict):
credentials: Optional[google.auth.credentials.Credentials] = None
project: Optional[str] = None
location: Optional[str] = None
class GoogleGenAIEmbedding(BaseEmbedding):
"""Google GenAI embeddings.
Args:
model_name (str): Model for embedding.
Defaults to "text-embedding-005".
api_key (Optional[str]): API key to access the model. Defaults to None.
embedding_config (Optional[types.EmbedContentConfigOrDict]): Embedding config to access the model. Defaults to None.
vertexai_config (Optional[VertexAIConfig]): Vertex AI config to access the model. Defaults to None.
http_options (Optional[types.HttpOptions]): HTTP options to access the model. Defaults to None.
debug_config (Optional[google.genai.client.DebugConfig]): Debug config to access the model. Defaults to None.
embed_batch_size (int): Batch size for embedding. Defaults to 100.
callback_manager (Optional[CallbackManager]): Callback manager to access the model. Defaults to None.
Examples:
`pip install llama-index-embeddings-google-genai`
```python
from llama_index.embeddings.google_genai import GoogleGenAIEmbedding
embed_model = GoogleGenAIEmbedding(model_name="text-embedding-005", api_key="...")
```
"""
_client: google.genai.Client = PrivateAttr()
_embedding_config: types.EmbedContentConfigOrDict = PrivateAttr()
embedding_config: Optional[types.EmbedContentConfigOrDict] = Field(
default=None, description="""Used to override embedding config."""
)
def __init__(
self,
model_name: str = "text-embedding-004",
api_key: Optional[str] = None,
embedding_config: Optional[types.EmbedContentConfigOrDict] = None,
vertexai_config: Optional[VertexAIConfig] = None,
http_options: Optional[types.HttpOptions] = None,
debug_config: Optional[google.genai.client.DebugConfig] = None,
embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
callback_manager: Optional[CallbackManager] = None,
**kwargs: Any,
):
super().__init__(
model_name=model_name,
embedding_config=embedding_config,
embed_batch_size=embed_batch_size,
callback_manager=callback_manager,
**kwargs,
)
# API keys are optional. The API can be authorised via OAuth (detected
# environmentally) or by the GOOGLE_API_KEY environment variable.
api_key = api_key or os.getenv("GOOGLE_API_KEY", None)
vertexai = vertexai_config is not None or os.getenv(
"GOOGLE_GENAI_USE_VERTEXAI", False
)
project = (vertexai_config or {}).get("project") or os.getenv(
"GOOGLE_CLOUD_PROJECT", None
)
location = (vertexai_config or {}).get("location") or os.getenv(
"GOOGLE_CLOUD_LOCATION", None
)
config_params: Dict[str, Any] = {
"api_key": api_key,
}
if vertexai_config is not None:
config_params.update(vertexai_config)
config_params["api_key"] = None
config_params["vertexai"] = True
elif vertexai:
config_params["project"] = project
config_params["location"] = location
config_params["api_key"] = None
config_params["vertexai"] = True
if http_options:
config_params["http_options"] = http_options
if debug_config:
config_params["debug_config"] = debug_config
self._client = google.genai.Client(**config_params)
@classmethod
def class_name(cls) -> str:
return "GeminiEmbedding"
def _embed_texts(
self, texts: List[str], task_type: Optional[str] = None
) -> List[List[float]]:
"""Embed texts."""
# Set the task type if it is not already set
if task_type and not self.embedding_config:
self.embedding_config = types.EmbedContentConfig(task_type=task_type)
results = self._client.models.embed_content(
model=self.model_name,
contents=texts,
config=self.embedding_config,
)
return [result.values for result in results.embeddings]
async def _aembed_texts(
self, texts: List[str], task_type: Optional[str] = None
) -> List[List[float]]:
"""Asynchronously embed texts."""
# Set the task type if it is not already set
if task_type and not self.embedding_config:
self.embedding_config = types.EmbedContentConfig(task_type=task_type)
results = await self._client.aio.models.embed_content(
model=self.model_name,
contents=texts,
config=self.embedding_config,
)
return [result.values for result in results.embeddings]
def _get_query_embedding(self, query: str) -> List[float]:
"""Get query embedding."""
return self._embed_texts([query], task_type="RETRIEVAL_QUERY")[0]
def _get_text_embedding(self, text: str) -> List[float]:
"""Get text embedding."""
return self._embed_texts([text], task_type="RETRIEVAL_DOCUMENT")[0]
def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
"""Get text embeddings."""
return self._embed_texts(texts, task_type="RETRIEVAL_DOCUMENT")
async def _aget_query_embedding(self, query: str) -> List[float]:
"""The asynchronous version of _get_query_embedding."""
return (await self._aembed_texts([query], task_type="RETRIEVAL_QUERY"))[0]
async def _aget_text_embedding(self, text: str) -> List[float]:
"""Asynchronously get text embedding."""
return (await self._aembed_texts([text], task_type="RETRIEVAL_DOCUMENT"))[0]
async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
"""Asynchronously get text embeddings."""
return await self._aembed_texts(texts, task_type="RETRIEVAL_DOCUMENT")
[build-system]
build-backend = "poetry.core.masonry.api"
requires = ["poetry-core"]
[tool.codespell]
check-filenames = true
check-hidden = true
skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
[tool.llamahub]
contains_example = false
import_path = "llama_index.embeddings.google_genai"
[tool.llamahub.class_authors]
GoogleGenAIEmbedding = "llama-index"
[tool.mypy]
disallow_untyped_defs = true
exclude = ["_static", "build", "examples", "notebooks", "venv"]
ignore_missing_imports = true
python_version = "3.8"
[tool.poetry]
authors = ["Your Name <you@example.com>"]
description = "llama-index embeddings google genai integration"
exclude = ["**/BUILD"]
license = "MIT"
name = "llama-index-embeddings-google-genai"
readme = "README.md"
version = "0.1.0"
[tool.poetry.dependencies]
python = ">=3.9,<4.0"
google-genai = ">=1.4.0"
llama-index-core = "^0.12.0"
[tool.poetry.group.dev.dependencies]
ipython = "8.10.0"
jupyter = "^1.0.0"
mypy = "0.991"
pre-commit = "3.2.0"
pylint = "2.15.10"
pytest = "7.2.1"
pytest-mock = "3.11.1"
ruff = "0.0.292"
tree-sitter-languages = "^1.8.0"
types-Deprecated = ">=0.1.0"
types-PyYAML = "^6.0.12.12"
types-protobuf = "^4.24.0.4"
types-redis = "4.5.5.0"
types-requests = "2.28.11.8"
types-setuptools = "67.1.0.0"
[tool.poetry.group.dev.dependencies.black]
extras = ["jupyter"]
version = "<=23.9.1,>=23.7.0"
[tool.poetry.group.dev.dependencies.codespell]
extras = ["toml"]
version = ">=v2.2.6"
[[tool.poetry.packages]]
include = "llama_index/"
python_tests(
interpreter_constraints=[">=3.10"],
dependencies=[
"llama-index-integrations/embeddings/llama-index-embeddings-google-genai:poetry",
],
)
from llama_index.core.base.embeddings.base import BaseEmbedding
from llama_index.embeddings.google_genai import GoogleGenAIEmbedding
def test_embedding_class():
emb = GoogleGenAIEmbedding(api_key="...")
assert isinstance(emb, BaseEmbedding)
# LlamaIndex Embeddings Integration: Google # LlamaIndex Embeddings Integration: Google
**NOTE:** This package is deprecated. Please use the `GoogleGenAIEmbedding` class instead, detailed [here](https://github.com/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/google_genai.ipynb).
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment