feat(llm): Support for Google Gemini LLMs and Embeddings (#1965)

* Support for Google Gemini LLMs and Embeddings Initial support for Gemini, enables usage of Google LLMs and embedding models (see settings-gemini.yaml) Install via poetry install --extras "llms-gemini embeddings-gemini" Notes: * had to bump llama-index-core to later version that supports Gemini * poetry --no-update did not work: Gemini/llama_index seem to require more (transient) updates to make it work... * fix: crash when gemini is not selected * docs: add gemini llm --------- Co-authored-by: Javier Martinez <[email protected]>
zylon-ai · Jul 8, 2024 · fc13368 · fc13368
1 parent 19a7c06
commit fc13368
Show file tree

Hide file tree

Showing 9 changed files with 382 additions and 59 deletions.
diff --git a/fern/docs/pages/manual/llms.mdx b/fern/docs/pages/manual/llms.mdx
@@ -199,3 +199,36 @@ Navigate to http://localhost:8001 to use the Gradio UI or to http://localhost:80
 For a fully private setup on Intel GPUs (such as a local PC with an iGPU, or discrete GPUs like Arc, Flex, and Max), you can use [IPEX-LLM](https://github.com/intel-analytics/ipex-llm).
 
 To deploy Ollama and pull models using IPEX-LLM, please refer to [this guide](https://ipex-llm.readthedocs.io/en/latest/doc/LLM/Quickstart/ollama_quickstart.html). Then, follow the same steps outlined in the [Using Ollama](#using-ollama) section to create a `settings-ollama.yaml` profile and run the private-GPT server.
+
+### Using Gemini
+
+If you cannot run a local model (because you don't have a GPU, for example) or for testing purposes, you may
+decide to run PrivateGPT using Gemini as the LLM and Embeddings model. In addition, you will benefit from
+multimodal inputs, such as text and images, in a very large contextual window.
+
+In order to do so, create a profile `settings-gemini.yaml` with the following contents:
+
+```yaml
+llm:
+  mode: gemini
+
+embedding:
+  mode: gemini
+
+gemini:
+  api_key: <your_gemini_api_key>                # You could skip this configuration and use the GEMINI_API_KEY env var instead
+  model: <gemini_model_to_use>                  # Optional model to use. Default is models/gemini-pro"
+  embedding_model: <gemini_embeddings_to_use>   # Optional model to use. Default is "models/embedding-001"
+```
+
+And run PrivateGPT loading that profile you just created:
+
+`PGPT_PROFILES=gemini make run`
+
+or
+
+`PGPT_PROFILES=gemini poetry run python -m private_gpt`
+
+When the server is started it will print a log *Application startup complete*.
+Navigate to http://localhost:8001 to use the Gradio UI or to http://localhost:8001/docs (API section) to try the API.
+
diff --git a/poetry.lock b/poetry.lock
diff --git a/private_gpt/components/embedding/embedding_component.py b/private_gpt/components/embedding/embedding_component.py
@@ -99,6 +99,20 @@ def __init__(self, settings: Settings) -> None:
                     azure_endpoint=azopenai_settings.azure_endpoint,
                     api_version=azopenai_settings.api_version,
                 )
+            case "gemini":
+                try:
+                    from llama_index.embeddings.gemini import (  # type: ignore
+                        GeminiEmbedding,
+                    )
+                except ImportError as e:
+                    raise ImportError(
+                        "Gemini dependencies not found, install with `poetry install --extras embeddings-gemini`"
+                    ) from e
+
+                self.embedding_model = GeminiEmbedding(
+                    api_key=settings.gemini.api_key,
+                    model_name=settings.gemini.embedding_model,
+                )
             case "mock":
                 # Not a random number, is the dimensionality used by
                 # the default embedding model

diff --git a/private_gpt/components/llm/llm_component.py b/private_gpt/components/llm/llm_component.py
@@ -190,5 +190,18 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                     azure_endpoint=azopenai_settings.azure_endpoint,
                     api_version=azopenai_settings.api_version,
                 )
+            case "gemini":
+                try:
+                    from llama_index.llms.gemini import (  # type: ignore
+                        Gemini,
+                    )
+                except ImportError as e:
+                    raise ImportError(
+                        "Google Gemini dependencies not found, install with `poetry install --extras llms-gemini`"
+                    ) from e
+                gemini_settings = settings.gemini
+                self.llm = Gemini(
+                    model_name=gemini_settings.model, api_key=gemini_settings.api_key
+                )
             case "mock":
                 self.llm = MockLLM()
diff --git a/private_gpt/settings/settings.py b/private_gpt/settings/settings.py
@@ -82,7 +82,14 @@ class DataSettings(BaseModel):
 
 class LLMSettings(BaseModel):
     mode: Literal[
-        "llamacpp", "openai", "openailike", "azopenai", "sagemaker", "mock", "ollama"
+        "llamacpp",
+        "openai",
+        "openailike",
+        "azopenai",
+        "sagemaker",
+        "mock",
+        "ollama",
+        "gemini",
     ]
     max_new_tokens: int = Field(
         256,
@@ -157,7 +164,9 @@ class HuggingFaceSettings(BaseModel):
 
 
 class EmbeddingSettings(BaseModel):
-    mode: Literal["huggingface", "openai", "azopenai", "sagemaker", "ollama", "mock"]
+    mode: Literal[
+        "huggingface", "openai", "azopenai", "sagemaker", "ollama", "mock", "gemini"
+    ]
     ingest_mode: Literal["simple", "batch", "parallel", "pipeline"] = Field(
         "simple",
         description=(
@@ -220,6 +229,18 @@ class OpenAISettings(BaseModel):
     )
 
 
+class GeminiSettings(BaseModel):
+    api_key: str
+    model: str = Field(
+        "models/gemini-pro",
+        description="Google Model to use. Example: 'models/gemini-pro'.",
+    )
+    embedding_model: str = Field(
+        "models/embedding-001",
+        description="Google Embedding Model to use. Example: 'models/embedding-001'.",
+    )
+
+
 class OllamaSettings(BaseModel):
     api_base: str = Field(
         "http://localhost:11434",
@@ -426,6 +447,7 @@ class Settings(BaseModel):
     huggingface: HuggingFaceSettings
     sagemaker: SagemakerSettings
     openai: OpenAISettings
+    gemini: GeminiSettings
     ollama: OllamaSettings
     azopenai: AzureOpenAISettings
     vectorstore: VectorstoreSettings

diff --git a/private_gpt/ui/ui.py b/private_gpt/ui/ui.py
@@ -444,6 +444,7 @@ def get_model_label() -> str | None:
                             "sagemaker": config_settings.sagemaker.llm_endpoint_name,
                             "mock": llm_mode,
                             "ollama": config_settings.ollama.llm_model,
+                            "gemini": config_settings.gemini.model,
                         }
 
                         if llm_mode not in model_mapping:

diff --git a/pyproject.toml b/pyproject.toml
@@ -24,10 +24,12 @@ llama-index-llms-openai = {version = "^0.1.25", optional = true}
 llama-index-llms-openai-like = {version ="^0.1.3", optional = true}
 llama-index-llms-ollama = {version ="^0.1.5", optional = true}
 llama-index-llms-azure-openai = {version ="^0.1.8", optional = true}
+llama-index-llms-gemini = {version ="^0.1.11", optional = true}
 llama-index-embeddings-ollama = {version ="^0.1.2", optional = true}
 llama-index-embeddings-huggingface = {version ="^0.2.2", optional = true}
 llama-index-embeddings-openai = {version ="^0.1.10", optional = true}
 llama-index-embeddings-azure-openai = {version ="^0.1.10", optional = true}
+llama-index-embeddings-gemini = {version ="^0.1.8", optional = true}
 llama-index-vector-stores-qdrant = {version ="^0.2.10", optional = true}
 llama-index-vector-stores-chroma = {version ="^0.1.10", optional = true}
 llama-index-vector-stores-postgres = {version ="^0.1.11", optional = true}
@@ -50,6 +52,9 @@ sentence-transformers = {version ="^3.0.1", optional = true}
 # Optional UI
 gradio = {version ="^4.37.2", optional = true}
 
+# Optional Google Gemini dependency
+google-generativeai = {version ="^0.5.4", optional = true}
+
 [tool.poetry.extras]
 ui = ["gradio"]
 llms-llama-cpp = ["llama-index-llms-llama-cpp"]
@@ -58,11 +63,13 @@ llms-openai-like = ["llama-index-llms-openai-like"]
 llms-ollama = ["llama-index-llms-ollama"]
 llms-sagemaker = ["boto3"]
 llms-azopenai = ["llama-index-llms-azure-openai"]
+llms-gemini = ["llama-index-llms-gemini", "google-generativeai"]
 embeddings-ollama = ["llama-index-embeddings-ollama"]
 embeddings-huggingface = ["llama-index-embeddings-huggingface"]
 embeddings-openai = ["llama-index-embeddings-openai"]
 embeddings-sagemaker = ["boto3"]
 embeddings-azopenai = ["llama-index-embeddings-azure-openai"]
+embeddings-gemini = ["llama-index-embeddings-gemini"]
 vector-stores-qdrant = ["llama-index-vector-stores-qdrant"]
 vector-stores-chroma = ["llama-index-vector-stores-chroma"]
 vector-stores-postgres = ["llama-index-vector-stores-postgres"]

diff --git a/settings-gemini.yaml b/settings-gemini.yaml
@@ -0,0 +1,10 @@
+llm:
+  mode: gemini
+
+embedding:
+  mode: gemini
+
+gemini:
+  api_key: ${GOOGLE_API_KEY:}
+  model: models/gemini-pro
+  embedding_model: models/embedding-001
diff --git a/settings.yaml b/settings.yaml
@@ -113,3 +113,8 @@ azopenai:
   api_version: "2023-05-15"
   embedding_model: text-embedding-ada-002
   llm_model: gpt-35-turbo
+
+gemini:
+  api_key: ${GOOGLE_API_KEY:}
+  model: models/gemini-pro
+  embedding_model: models/embedding-001